feat: add write_tensor

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

.github/workflows/build.yml

@@ -26,10 +26,8 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
-      - name: Use beta toolchain
-        run: |
-          rustup toolchain install beta --component rustfmt clippy
-          rustup default beta
+      - name: Use latest toolchain
+        run: rustup update
 
       - name: Check format
         run: cargo fmt --check

CHANGELOG.md

@@ -5,12 +5,18 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.2.1] - 2025-03-28
+
+### Added
+
+- Add `write_array` to write Nd array to `fmt::Formatter`;
+
 ## [0.2.0] - 2025-02-23
 
 ### Added
 
-- Add `num_elements` to calculate the number of elements in array.
-- Add `element_offset` to calculate the offset of element at the given index.
+- Add `num_elements` to calculate the number of elements in array;
+- Add `element_offset` to calculate the offset of element at the given index;
 
 ### Changed
 
@@ -20,4 +26,5 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fix merging dims with length 1;
 
+[0.2.1]: https://github.com/InfiniTensor/ndarray-layout/compare/v0.2.0...v0.2.1
 [0.2.0]: https://github.com/InfiniTensor/ndarray-layout/releases/tag/v0.2.0

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "ndarray-layout"
 description = "This crate provides definitions and transformations for multi-dimensional array data layouts."
-version = "0.2.0"
+version = "0.2.1"
 edition = "2024"
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 repository = "https://github.com/InfiniTensor/ndarray-layout.git"

src/fmt.rs

@@ -0,0 +1,111 @@
+use crate::ArrayLayout;
+use std::fmt;
+
+impl<const N: usize> ArrayLayout<N> {
+    /// 高维数组格式化。
+    pub fn write_array<T: fmt::Display>(
+        &self,
+        f: &mut fmt::Formatter,
+        ptr: *const T,
+    ) -> fmt::Result {
+        match self.ndim() {
+            0 => {
+                write!(f, "array<> = [{}]", unsafe {
+                    ptr.byte_offset(self.offset()).read_unaligned()
+                })
+            }
+            1 => {
+                let &[n] = self.shape() else { unreachable!() };
+                let &[s] = self.strides() else { unreachable!() };
+
+                writeln!(f, "array<{n}>[")?;
+                let ptr = unsafe { ptr.byte_offset(self.offset()) };
+                for i in 0..n as isize {
+                    writeln!(f, "    {}", unsafe {
+                        ptr.byte_offset(i * s).read_unaligned()
+                    })?
+                }
+                writeln!(f, "]")?;
+                Ok(())
+            }
+            _ => {
+                let mut title = "array<".to_string();
+                for d in self.shape() {
+                    title.push_str(&format!("{d}x"))
+                }
+                assert_eq!(title.pop(), Some('x'));
+                title.push('>');
+
+                let mut stack = Vec::with_capacity(self.ndim() - 2);
+                self.write_recursive(f, ptr, &title, &mut stack)
+            }
+        }
+    }
+
+    fn write_recursive<T: fmt::Display>(
+        &self,
+        f: &mut fmt::Formatter,
+        ptr: *const T,
+        title: &str,
+        indices: &mut Vec<usize>,
+    ) -> fmt::Result {
+        match *self.shape() {
+            [] | [_] => unreachable!(),
+            [rows, cols] => {
+                write!(f, "{title}[")?;
+                for i in indices {
+                    write!(f, "{i}, ")?
+                }
+                writeln!(f, "..]")?;
+
+                let &[rs, cs] = self.strides() else {
+                    unreachable!()
+                };
+
+                let ptr = unsafe { ptr.byte_offset(self.offset()) };
+                for r in 0..rows as isize {
+                    for c in 0..cols as isize {
+                        write!(f, "{} ", unsafe {
+                            ptr.byte_offset(r * rs + c * cs).read_unaligned()
+                        })?
+                    }
+                    writeln!(f)?
+                }
+            }
+            [batch, ..] => {
+                for i in 0..batch {
+                    indices.push(i);
+                    self.index(0, i).write_recursive(f, ptr, title, indices)?;
+                    indices.pop();
+                }
+            }
+        }
+        Ok(())
+    }
+}
+
+#[test]
+fn test() {
+    const DATA: &[u8] = &[1, 2, 3, 4, 5, 6, 7, 8, 9, 0];
+
+    struct Tensor(ArrayLayout<4>);
+
+    impl fmt::Display for Tensor {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            self.0.write_array(f, DATA.as_ptr())
+        }
+    }
+
+    let tensor = Tensor(ArrayLayout::<4>::new_contiguous(
+        &[DATA.len()],
+        crate::Endian::BigEndian,
+        1,
+    ));
+    println!("{}", tensor);
+
+    let tensor = Tensor(tensor.0.tile_be(0, &[1, DATA.len()]).broadcast(0, 6));
+    println!("{}", tensor);
+
+    let tensor = Tensor(tensor.0.tile_be(0, &[2, 3]).tile_be(2, &[5, 2]));
+    println!("{}", tensor);
+}

src/lib.rs

@@ -184,6 +184,7 @@ impl<const N: usize> ArrayLayout<N> {
     }
 }
 
+mod fmt;
 mod transform;
 pub use transform::{BroadcastArg, IndexArg, MergeArg, SliceArg, Split, TileArg};