Cargo.lock

@@ -302,6 +302,7 @@ dependencies = [
  "ndarray",
  "num",
  "rand",
+ "rayon",
  "serde",
  "serde_json",
  "static_assertions",

Cargo.toml

@@ -19,6 +19,7 @@ bytemuck = "1.14.0"
 getset = "0.1.2"
 itertools = "0.12.0"
 num = "0.4.1"
+rayon = "1.8.0"
 serde = { version = "1.0.193", features = ["derive"] }
 thiserror = "1.0.52"
 

benches/manifold_benchmark.rs

@@ -3,9 +3,11 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use manifold::*;
 use rand::Rng;
 
+const SIZE: usize = 10000;
+
 fn random_tensor_r2_manifold() -> Tensor<f64, 2> {
     let mut rng = rand::thread_rng();
-    let mut tensor = tensor!([[0.0; 1000]; 1000]);
+    let mut tensor = tensor!([[0.0; SIZE]; SIZE]);
     for i in 0..tensor.len() {
         tensor[i] = rng.gen();
     }
@@ -14,7 +16,7 @@ fn random_tensor_r2_manifold() -> Tensor<f64, 2> {
 
 fn random_tensor_r2_ndarray() -> ndarray::Array2<f64> {
     let mut rng = rand::thread_rng();
-    let (rows, cols) = (1000, 1000);
+    let (rows, cols) = (SIZE, SIZE);
     let mut tensor = ndarray::Array2::<f64>::zeros((rows, cols));
     for i in 0..rows {
         for j in 0..cols {
@@ -40,7 +42,7 @@ fn tensor_product(c: &mut Criterion) {
                     let a = random_tensor_r2_manifold();
                     let b = random_tensor_r2_manifold();
                     let c = a + b;
-                    assert!(c.shape().as_array() == &[1000, 1000]);
+                    assert!(c.shape().as_array() == &[SIZE, SIZE]);
                 })
             },
         );
@@ -53,7 +55,7 @@ fn tensor_product(c: &mut Criterion) {
                     let a = random_tensor_r2_ndarray();
                     let b = random_tensor_r2_ndarray();
                     let c = a + b;
-                    assert!(c.shape() == &[1000, 1000]);
+                    assert!(c.shape() == &[SIZE, SIZE]);
                 })
             },
         );

src/axis.rs

@@ -1,118 +0,0 @@
-use super::*;
-use getset::{Getters, MutGetters};
-
-#[derive(Clone, Debug, Getters)]
-pub struct TensorAxis<'a, T: Value, const R: usize> {
-    #[getset(get = "pub")]
-    tensor: &'a Tensor<T, R>,
-    #[getset(get = "pub")]
-    dim: usize,
-}
-
-impl<'a, T: Value, const R: usize> TensorAxis<'a, T, R> {
-    pub fn new(tensor: &'a Tensor<T, R>, dim: usize) -> Self {
-        assert!(dim < R, "TensorAxis out of bounds");
-        Self { tensor, dim }
-    }
-
-    pub fn len(&self) -> usize {
-        self.tensor.shape().get(self.dim)
-    }
-
-    pub fn shape(&self) -> &TensorShape<R> {
-        self.tensor.shape()
-    }
-
-    pub fn iter_level(&'a self, level: usize) -> TensorAxisIterator<'a, T, R> {
-        assert!(level < self.len(), "Level out of bounds");
-        let mut index = TensorIndex::new(self.shape().clone(), [0; R]);
-        index.set_axis(self.dim, level);
-        TensorAxisIterator::new(self)
-            .set_start(level)
-            .set_end(level + 1)
-    }
-}
-
-#[derive(Clone, Debug, Getters, MutGetters)]
-pub struct TensorAxisIterator<'a, T: Value, const R: usize> {
-    #[getset(get = "pub")]
-    axis: &'a TensorAxis<'a, T, R>,
-    #[getset(get = "pub", get_mut = "pub")]
-    index: TensorIndex<R>,
-    #[getset(get = "pub")]
-    end: Option<usize>,
-}
-
-impl<'a, T: Value, const R: usize> TensorAxisIterator<'a, T, R> {
-    pub fn new(axis: &'a TensorAxis<'a, T, R>) -> Self {
-        Self {
-            axis,
-            index: TensorIndex::new(axis.shape().clone(), [0; R]),
-            end: None,
-        }
-    }
-
-    pub fn set_start(self, start: usize) -> Self {
-        assert!(start < self.axis().len(), "Start out of bounds");
-        let mut index = TensorIndex::new(self.axis().shape().clone(), [0; R]);
-        index.set_axis(self.axis.dim, start);
-        Self {
-            axis: self.axis(),
-            index,
-            end: None,
-        }
-    }
-
-    pub fn set_end(self, end: usize) -> Self {
-        assert!(end <= self.axis().len(), "End out of bounds");
-        Self {
-            axis: self.axis(),
-            index: self.index().clone(),
-            end: Some(end),
-        }
-    }
-
-    pub fn set_level(self, level: usize) -> Self {
-        assert!(level < self.axis().len(), "Level out of bounds");
-        self.set_start(level).set_end(level + 1)
-    }
-
-    pub fn level(&'a self, level: usize) -> impl Iterator<Item = &'a T> + 'a {
-        Self::new(self.axis()).set_level(level)
-    }
-
-    pub fn axis_max_idx(&self) -> usize {
-        self.end().unwrap_or(self.axis().len())
-    }
-
-    pub fn axis_idx(&self) -> usize {
-        self.index().get_axis(*self.axis().dim())
-    }
-
-    pub fn axis_dim(&self) -> usize {
-        self.axis().dim().clone()
-    }
-}
-
-impl<'a, T: Value, const R: usize> Iterator for TensorAxisIterator<'a, T, R> {
-    type Item = &'a T;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.axis_idx() == self.axis_max_idx() {
-            return None;
-        }
-        let result = unsafe { self.axis().tensor().get_unchecked(self.index) };
-        let axis_dim = self.axis_dim();
-        self.index_mut().inc_axis(axis_dim);
-        Some(result)
-    }
-}
-
-impl<'a, T: Value, const R: usize> IntoIterator for &'a TensorAxis<'a, T, R> {
-    type Item = &'a T;
-    type IntoIter = TensorAxisIterator<'a, T, R>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        TensorAxisIterator::new(&self)
-    }
-}

src/index.rs

@@ -2,9 +2,18 @@ use super::*;
 use getset::{Getters, MutGetters};
 use std::{
     cmp::Ordering,
-    ops::{Add, Index, IndexMut, Sub},
+    ops::{Index, IndexMut},
 };
 
+/// A multi-dimensional index into a tensor.
+///
+/// ```
+/// use manifold::*;
+///
+/// let shape = TensorShape::new([2, 3, 4]);
+/// let index = TensorIndex::new(shape, [1, 2, 3]);
+/// assert_eq!(index.flat(), 23);
+/// ```
 #[derive(Clone, Copy, Debug, Getters, MutGetters)]
 pub struct TensorIndex<const R: usize> {
     #[getset(get = "pub", get_mut = "pub")]
@@ -16,13 +25,30 @@ pub struct TensorIndex<const R: usize> {
 // ---- Construction and Initialization ---------------------------------------
 
 impl<const R: usize> TensorIndex<R> {
+    /// Creates a new `TensorIndex` instance.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let index = TensorIndex::new(shape, [1, 2, 3]);
+    /// assert_eq!(index.flat(), 23);
+    /// ```
     pub fn new(shape: TensorShape<R>, indices: [usize; R]) -> Self {
-        if !shape.check_indices(indices) {
+        let index = Self { indices, shape };
-            panic!("indices out of bounds");
+        assert!(index.check_indices(indices));
-        }
+        index
-        Self { indices, shape }
     }
 
+    /// Creates a new `TensorIndex` instance with all indices set to zero.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let index = TensorIndex::zero(shape);
+    /// assert_eq!(index.flat(), 0);
+    /// ```
     pub const fn zero(shape: TensorShape<R>) -> Self {
         Self {
             indices: [0; R],
@@ -30,106 +56,185 @@ impl<const R: usize> TensorIndex<R> {
         }
     }
 
-    pub fn last(shape: TensorShape<R>) -> Self {
+    pub fn from_flat(shape: TensorShape<R>, flat_index: usize) -> Self {
-        let max_indices =
+        let mut indices = [0; R];
-            shape.as_array().map(|dim_size| dim_size.saturating_sub(1));
+        let mut remaining = flat_index;
-        Self {
+
-            indices: max_indices,
+        // - The method iterates over the dimensions of the tensor in reverse
-            shape,
+        //   order (assuming row-major order).
-        }
+        // - In each iteration, it uses the modulo operation to find the index
+        //   in the current dimension and integer division to reduce the flat
+        //   index for the next higher dimension.
+        // - This process is repeated for each dimension to build the
+        //   multi-dimensional index.
+        for (idx, &dim_size) in indices.iter_mut().zip(shape.0.iter()).rev() {
+            *idx = remaining % dim_size;
+            remaining /= dim_size;
+        }
+
+        // Reverse the indices to match the original dimension order
+        indices.reverse();
+        Self::new(shape.clone(), indices)
     }
 }
 
+// ---- Trivial Functions -----------------------------------------------------
+
 impl<const R: usize> TensorIndex<R> {
-    pub fn is_zero(&self) -> bool {
+    /// Returns `true` if all indices are zero.
-        self.indices.iter().all(|&i| i == 0)
-    }
-
-    pub fn is_overflow(&self) -> bool {
-        // Check if the last index is equal to the size of the last dimension
-        self.indices[0] >= self.shape.get(R - 1)
-    }
-
-    pub fn reset(&mut self) {
-        self.indices = [0; R];
-    }
-
-    /// Increments the index and returns a boolean indicating whether the end
-    /// has been reached.
     ///
-    /// # Returns
+    /// ```
-    /// `true` if the increment does not overflow and is still within bounds;
+    /// use manifold::*;
-    /// `false` if it overflows, indicating the end of the tensor.
+    ///
-    pub fn inc(&mut self) -> bool {
+    /// let shape = TensorShape::new([2, 3, 4]);
-        if self.indices()[0] >= self.shape().get(0) {
+    /// let index = TensorIndex::zero(shape);
-            return false;
+    /// assert!(index.is_zero());
+    /// ```
+    pub fn is_zero(&self) -> bool {
+        self.indices().iter().all(|&i| i == 0)
     }
+
+    /// Returns `true` if the last index is equal to the size of the last
+    /// dimension.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([1, 1]);
+    /// let mut index = TensorIndex::zero(shape);
+    /// assert!(!index.is_max());
+    /// index.inc();
+    /// assert!(index.is_max());
+    /// ```
+    pub fn is_max(&self) -> bool {
+        self.indices().get(0) == self.shape().dim_size(0)
+    }
+}
+
+// ---- Utils -----------------------------------------------------------------
+
+impl<const R: usize> TensorIndex<R> {
+    /// Resets the index to zero.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let mut index = TensorIndex::new(shape, [1, 2, 3]);
+    /// assert_eq!(index.flat(), 23);
+    /// index.reset();
+    /// assert_eq!(index.flat(), 0);
+    /// ```
+    pub fn reset(&mut self) {
+        *self.indices_mut() = [0; R];
+    }
+
+    fn check_indices(&self, indices: [usize; R]) -> bool {
+        indices
+            .iter()
+            .zip(self.shape().as_array().iter())
+            .all(|(&idx, &dim_size)| idx < dim_size)
+    }
+}
+
+// ---- Increment and Decrement -----------------------------------------------
+
+impl<const R: usize> TensorIndex<R> {
+    /// Increments the index by one.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let mut index = TensorIndex::zero(shape);
+    /// assert_eq!(index.flat(), 0);
+    /// index.inc();
+    /// assert_eq!(index.flat(), 1);
+    /// ```
+    pub fn inc(&mut self) {
+        if self.indices().get(0) == self.shape().dim_size(0) {
+            return;
+        }
+
         let shape = self.shape().as_array().clone();
-        let mut carry = 1;
-        for (i, &dim_size) in self.indices.iter_mut().zip(&shape).rev() {
-            if carry == 1 {
-                *i += 1;
-                if *i >= dim_size {
-                    *i = 0; // Reset index in this dimension and carry over
-                } else {
-                    carry = 0; // Increment successful, no carry needed
-                }
-            }
-        }
 
-        // If carry is still 1 after the loop, it means we've incremented past
+        for (i, dim) in self.indices_mut().iter_mut().rev().enumerate() {
-        // the last dimension
+            *dim += 1;
-        if carry == 1 {
+            if *dim < shape[i] {
-            // Set the index to an invalid state to indicate the end of the
-            // iteration indicated by setting the first index to the
-            // size of the first dimension
-            self.indices[0] = self.shape.as_array()[0];
-            return true; // Indicate that the iteration is complete
-        }
-        false
-    }
-
-    // fn inc_axis
-
-    pub fn inc_axis(&mut self, fixed_axis: usize) {
-        assert!(fixed_axis < R, "TensorAxis out of bounds");
-        assert!(
-            self.indices()[fixed_axis] < self.shape().get(fixed_axis),
-            "Index out of bounds"
-        );
-
-        // Try to increment non-fixed axes
-        for i in (0..R).rev() {
-            if i != fixed_axis {
-                if self.indices[i] + 1 < self.shape.get(i) {
-                    self.indices[i] += 1;
-                    return;
-                } else {
-                    self.indices[i] = 0;
-                }
-            }
-        }
-
-        if self.indices[fixed_axis] < self.shape.get(fixed_axis) {
-            self.indices[fixed_axis] += 1;
-            for i in 0..R {
-                if i != fixed_axis {
-                    self.indices[i] = 0;
-                }
-            }
                 return;
             }
+            *dim = 0;
         }
 
-    pub fn inc_transposed(&mut self, order: &[usize; R]) -> bool {
+        self.indices_mut()[0] = *self.shape().dim_size(0).unwrap();
-        if self.indices()[order[0]] >= self.shape().get(order[0]) {
+    }
-            return false;
+
+    /// Increments the index by one, with the specified axis fixed,
+    /// i.e. the index of the fixed axis is incremented only if the
+    /// index of the other axes reaches the maximum.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let mut index = TensorIndex::zero(shape);
+    /// assert_eq!(index.flat(), 0);
+    /// index.inc_fixed_axis(Axis(0));
+    /// assert_eq!(index.flat(), 4);
+    /// index.inc_fixed_axis(Axis(1));
+    /// assert_eq!(index.flat(), 5);
+    /// ```
+    pub fn inc_fixed_axis(&mut self, Axis(ax): Axis) {
+        let shape = self.shape().as_array().clone();
+        assert!(ax < R, "TensorAxis out of bounds");
+        if self.indices().get(ax) == self.shape().dim_size(ax) {
+            return;
+        }
+
+        // Iterate over all axes, skipping the fixed axis 'ax'
+        for (i, dim) in self.indices_mut().iter_mut().rev().enumerate() {
+            if i != ax {
+                // Skip the fixed axis
+                *dim += 1;
+                if *dim < shape[i] {
+                    return; // No carry over needed
+                }
+                *dim = 0; // Reset the current axis and carry over to the next
+            }
+        }
+
+        // Handle the case where incrementing has reached the end
+        if self.indices().get(ax) < self.shape().dim_size(ax) {
+            self.indices_mut()[ax] += 1;
+        } else {
+            // Reset if the fixed axis also overflows
+            self.indices_mut()[ax] = *self.shape().dim_size(ax).unwrap();
+        }
+    }
+
+    /// Increments the index by one, reordering the order in which the
+    /// axes are incremented.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let mut index = TensorIndex::zero(shape);
+    /// assert_eq!(index.flat(), 0);
+    /// index.inc_transposed(&[2, 1, 0]);
+    /// assert_eq!(index.flat(), 12);
+    /// index.inc_transposed(&[0, 1, 2]);
+    /// assert_eq!(index.flat(), 13);
+    /// ```
+    pub fn inc_transposed(&mut self, order: &[usize; R]) {
+        if self.indices().get(order[0]) == self.shape().dim_size(order[0]) {
+            return;
         }
 
         let mut carry = 1;
 
         for i in order.iter().rev() {
-            let dim_size = self.shape().get(*i);
+            let dim_size = self.shape().dim_size(*i).unwrap().clone();
             let i = self.index_mut(*i);
             if carry == 1 {
                 *i += 1;
@@ -142,13 +247,21 @@ impl<const R: usize> TensorIndex<R> {
         }
 
         if carry == 1 {
-            self.indices_mut()[order[0]] = self.shape().get(order[0]);
+            self.indices_mut()[order[0]] = *self.shape().dim_size(order[0]).unwrap();
-            return true;
+        }
-        }
-
-        false
     }
 
+    /// Decrements the index by one.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let mut index = TensorIndex::new(shape, [1, 2, 3]);
+    /// assert_eq!(index.flat(), 23);
+    /// index.dec();
+    /// assert_eq!(index.flat(), 22);
+    /// ```
     pub fn dec(&mut self) {
         // Check if already at the start
         if self.indices.iter().all(|&i| i == 0) {
@@ -169,102 +282,22 @@ impl<const R: usize> TensorIndex<R> {
             }
         }
     }
-
-    pub fn dec_axis(&mut self, fixed_axis: usize) -> bool {
-        // Check if the fixed axis index is already in an invalid state
-        if self.indices[fixed_axis] == self.shape.get(fixed_axis) {
-            return false;
 }
 
-        // Try to decrement non-fixed axes
+// ---- Conversion to Flat Index ----------------------------------------------
-        for i in (0..R).rev() {
-            if i != fixed_axis {
-                if self.indices[i] > 0 {
-                    self.indices[i] -= 1;
-                    return true;
-                } else {
-                    self.indices[i] = self.shape.get(i) - 1;
-                }
-            }
-        }
-
-        // Decrement the fixed axis if possible and reset other axes to their
-        // max
-        if self.indices[fixed_axis] > 0 {
-            self.indices[fixed_axis] -= 1;
-            for i in 0..R {
-                if i != fixed_axis {
-                    self.indices[i] = self.shape.get(i) - 1;
-                }
-            }
-        } else {
-            // Fixed axis already at minimum, set to invalid state
-            self.indices[fixed_axis] = self.shape.get(fixed_axis);
-        }
-
-        true
-    }
-
-    pub fn dec_transposed(&mut self, order: [usize; R]) {
-        // Iterate over the axes in the specified order
-        for &axis in &order {
-            // Try to decrement the current axis
-            if self.indices[axis] > 0 {
-                self.indices[axis] -= 1;
-                // Reset all preceding axes in the order to their maximum
-                for &prev_axis in &order {
-                    if prev_axis == axis {
-                        break;
-                    }
-                    self.indices[prev_axis] = self.shape.get(prev_axis) - 1;
-                }
-                return;
-            }
-        }
-
-        // If no axis can be decremented, set the first axis in the order to
-        // indicate overflow
-        self.indices[order[0]] = self.shape.get(order[0]);
-    }
 
+impl<const R: usize> TensorIndex<R> {
     /// Converts the multi-dimensional index to a flat index.
     ///
-    /// This method calculates the flat index corresponding to the
+    /// # Examples
-    /// multi-dimensional index stored in `self.indices`, given the shape of
-    /// the tensor stored in `self.shape`. The calculation is based on the
-    /// assumption that the tensor is stored in row-major order,
-    /// where the last dimension varies the fastest.
     ///
-    /// # Returns
+    /// ```
-    /// The flat index corresponding to the multi-dimensional index.
+    /// use manifold::*;
     ///
-    /// # How It Works
+    /// let shape = TensorShape::new([2, 3, 4]);
-    /// - The method iterates over each pair of corresponding index and
+    /// let index = TensorIndex::new(shape, [1, 2, 3]);
-    ///   dimension size, starting from the last dimension (hence `rev()` is
+    /// assert_eq!(index.flat(), 23);
-    ///   used for reverse iteration).
+    /// ```
-    /// - In each iteration, it performs two main operations:
-    ///   1. **Index Contribution**: Multiplies the current index (`idx`) by a
-    ///      running product of dimension sizes (`product`). This calculates the
-    ///      contribution of the current index to the overall flat index.
-    ///   2. **Product Update**: Multiplies `product` by the current dimension
-    ///      size (`dim_size`). This updates `product` for the next iteration,
-    ///      as each dimension contributes to the flat index based on the sizes
-    ///      of all preceding dimensions.
-    /// - The `fold` operation accumulates these results, starting with an
-    ///   initial state of `(0, 1)` where `0` is the initial flat index and `1`
-    ///   is the initial product.
-    /// - The final flat index is obtained after the last iteration, which is
-    ///   the first element of the tuple resulting from the `fold`.
-    ///
-    /// # Example
-    /// Consider a tensor with shape `[3, 4, 5]` and an index `[1, 2, 3]`.
-    /// - Starting with a flat index of 0 and a product of 1,
-    /// - For the last dimension (size 5), add 3 * 1 to the flat index. Update
-    ///   the product to 1 * 5 = 5.
-    /// - For the second dimension (size 4), add 2 * 5 to the flat index. Update
-    ///   the product to 5 * 4 = 20.
-    /// - For the first dimension (size 3), add 1 * 20 to the flat index. The
-    ///   final flat index is 3 + 10 + 20 = 33.
     pub fn flat(&self) -> usize {
         self.indices()
             .iter()
@@ -275,37 +308,9 @@ impl<const R: usize> TensorIndex<R> {
             })
             .0
     }
-
-    pub fn set_axis(&mut self, axis: usize, value: usize) {
-        assert!(axis < R, "TensorAxis out of bounds");
-        // assert!(value < self.shape.get(axis), "Value out of bounds");
-        self.indices[axis] = value;
 }
 
-    pub fn try_set_axis(&mut self, axis: usize, value: usize) -> bool {
+// --- Equality and Ordering --------------------------------------------------
-        assert!(axis < R, "TensorAxis out of bounds");
-        if value < self.shape.get(axis) {
-            self.indices[axis] = value;
-            true
-        } else {
-            false
-        }
-    }
-
-    pub fn get_axis(&self, axis: usize) -> usize {
-        assert!(axis < R, "TensorAxis out of bounds");
-        self.indices[axis]
-    }
-
-    pub fn iter_transposed(
-        &self,
-        order: [usize; R],
-    ) -> TensorIndexTransposedIterator<R> {
-        TensorIndexTransposedIterator::new(self.shape().clone(), order)
-    }
-}
-
-// --- blanket impls ---
 
 impl<const R: usize> PartialEq for TensorIndex<R> {
     fn eq(&self, other: &Self) -> bool {
@@ -327,6 +332,8 @@ impl<const R: usize> Ord for TensorIndex<R> {
     }
 }
 
+// ---- Indexing --------------------------------------------------------------
+
 impl<const R: usize> Index<usize> for TensorIndex<R> {
     type Output = usize;
 
@@ -341,17 +348,18 @@ impl<const R: usize> IndexMut<usize> for TensorIndex<R> {
     }
 }
 
+// ---- Conversion to TensorIndex ---------------------------------------------
+
 impl<const R: usize> From<(TensorShape<R>, [usize; R])> for TensorIndex<R> {
     fn from((shape, indices): (TensorShape<R>, [usize; R])) -> Self {
-        assert!(shape.check_indices(indices));
+        // assert!(shape.check_indices(indices));
         Self::new(shape, indices)
     }
 }
 
 impl<const R: usize> From<(TensorShape<R>, usize)> for TensorIndex<R> {
     fn from((shape, flat_index): (TensorShape<R>, usize)) -> Self {
-        let indices = shape.index_from_flat(flat_index).indices;
+        Self::from_flat(shape, flat_index)
-        Self::new(shape, indices)
     }
 }
 
@@ -367,6 +375,8 @@ impl<T: Value, const R: usize> From<Tensor<T, R>> for TensorIndex<R> {
     }
 }
 
+// ---- Display ---------------------------------------------------------------
+
 impl<const R: usize> std::fmt::Display for TensorIndex<R> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "[")?;
@@ -384,113 +394,3 @@ impl<const R: usize> std::fmt::Display for TensorIndex<R> {
         write!(f, "]")
     }
 }
-
-// ---- Arithmetic Operations ----
-
-impl<const R: usize> Add for TensorIndex<R> {
-    type Output = Self;
-
-    fn add(self, rhs: Self) -> Self::Output {
-        assert_eq!(self.shape, rhs.shape, "TensorShape mismatch");
-
-        let mut result_indices = [0; R];
-        for i in 0..R {
-            result_indices[i] = self.indices[i] + rhs.indices[i];
-        }
-
-        Self {
-            indices: result_indices,
-            shape: self.shape,
-        }
-    }
-}
-
-impl<const R: usize> Sub for TensorIndex<R> {
-    type Output = Self;
-
-    fn sub(self, rhs: Self) -> Self::Output {
-        assert_eq!(self.shape, rhs.shape, "TensorShape mismatch");
-
-        let mut result_indices = [0; R];
-        for i in 0..R {
-            result_indices[i] = self.indices[i].saturating_sub(rhs.indices[i]);
-        }
-
-        Self {
-            indices: result_indices,
-            shape: self.shape,
-        }
-    }
-}
-
-// ---- Iterator ----
-
-pub struct TensorIndexIterator<const R: usize> {
-    current: TensorIndex<R>,
-    end: bool,
-}
-
-impl<const R: usize> TensorIndexIterator<R> {
-    pub fn new(shape: TensorShape<R>) -> Self {
-        Self {
-            current: TensorIndex::zero(shape),
-            end: false,
-        }
-    }
-}
-
-impl<const R: usize> Iterator for TensorIndexIterator<R> {
-    type Item = TensorIndex<R>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.end {
-            return None;
-        }
-
-        let result = self.current;
-        self.end = self.current.inc();
-        Some(result)
-    }
-}
-
-impl<const R: usize> IntoIterator for TensorIndex<R> {
-    type Item = TensorIndex<R>;
-    type IntoIter = TensorIndexIterator<R>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        TensorIndexIterator {
-            current: self,
-            end: false,
-        }
-    }
-}
-
-pub struct TensorIndexTransposedIterator<const R: usize> {
-    current: TensorIndex<R>,
-    order: [usize; R],
-    end: bool,
-}
-
-impl<const R: usize> TensorIndexTransposedIterator<R> {
-    pub fn new(shape: TensorShape<R>, order: [usize; R]) -> Self {
-        Self {
-            current: TensorIndex::zero(shape),
-            end: false,
-            order,
-        }
-    }
-}
-
-impl<const R: usize> Iterator for TensorIndexTransposedIterator<R> {
-    type Item = TensorIndex<R>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.end {
-            return None;
-        }
-
-        let result = self.current;
-        self.end = self.current.inc_transposed(&self.order);
-        Some(result)
-    }
-}

src/iterators.rs

@@ -0,0 +1,71 @@
+use super::*;
+use std::fmt::{Display, Formatter, Result as FmtResult};
+
+// ---- Iterator --------------------------------------------------------------
+
+pub struct TensorIterator<'a, T: Value, const R: usize> {
+    tensor: &'a Tensor<T, R>,
+    index: TensorIndex<R>,
+}
+
+impl<'a, T: Value, const R: usize> TensorIterator<'a, T, R> {
+    pub fn new(tensor: &'a Tensor<T, R>) -> Self {
+        Self {
+            tensor,
+            index: tensor.shape().index_zero(),
+        }
+    }
+}
+
+impl<'a, T: Value, const R: usize> Iterator for TensorIterator<'a, T, R> {
+    type Item = &'a T;
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.index.is_max() {
+            return None;
+        }
+
+        let result = unsafe { self.tensor.get_unchecked(self.index) };
+        self.index.inc();
+        Some(result)
+    }
+}
+
+impl<'a, T: Value, const R: usize> IntoIterator for &'a Tensor<T, R> {
+    type Item = &'a T;
+    type IntoIter = TensorIterator<'a, T, R>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        TensorIterator::new(self)
+    }
+}
+
+impl<'a, T: Value, const R: usize> Display for TensorIterator<'a, T, R> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
+        // Print the current index and flat index
+        write!(
+            f,
+            "Current Index: {}, Flat Index: {}",
+            self.index,
+            self.index.flat()
+        )?;
+
+        // Print the tensor elements, highlighting the current element
+        write!(f, ", Tensor Elements: [")?;
+        for (i, elem) in self.tensor.buffer().iter().enumerate() {
+            if i == self.index.flat() {
+                write!(f, "*{}*", elem)?; // Highlight the current element
+            } else {
+                write!(f, "{}", elem)?;
+            }
+            if i < self.tensor.buffer().len() - 1 {
+                write!(f, ", ")?;
+            }
+        }
+
+        write!(f, "]")
+    }
+}
+
+// ---- Axis Iterator ---------------------------------------------------------
+
+// ---- Transposed Iterator ---------------------------------------------------

src/lib.rs

@@ -2,14 +2,36 @@
 #![feature(generic_const_exprs)]
 #![warn(clippy::all)]
 
-pub mod axis;
 pub mod error;
 pub mod index;
+pub mod iterators;
 pub mod shape;
 pub mod tensor;
-pub mod value;
 
-pub use {axis::*, error::*, index::*, shape::*, tensor::*, value::*};
+pub use {error::*, index::*, iterators::*, shape::*, tensor::*};
+
+use num::{Num, One, Zero};
+use serde::{Deserialize, Serialize};
+use std::{fmt::Display, iter::Sum};
+
+/// A trait for types that can be used as values in a tensor.
+pub trait Value:
+    Num + Zero + One + Copy + Clone + Display + Serialize + Deserialize<'static>
+{
+}
+
+impl<T> Value for T where
+    T: Num
+        + Zero
+        + One
+        + Copy
+        + Clone
+        + Display
+        + Serialize
+        + Deserialize<'static>
+        + Sum
+{
+}
 
 #[macro_export]
 macro_rules! tensor {

src/shape.rs

@@ -4,18 +4,90 @@ use serde::de::{self, Deserialize, Deserializer, SeqAccess, Visitor};
 use serde::ser::{Serialize, SerializeTuple, Serializer};
 use std::fmt::{Formatter, Result as FmtResult};
 
+/// A tensor's shape.
+///
+/// ```
+/// use manifold::*;
+///
+/// let shape = shape!([2, 3]);
+/// assert_eq!(shape.dim_size(0), Some(&2));
+/// assert_eq!(shape.dim_size(1), Some(&3));
+/// assert_eq!(shape.dim_size(2), None);
+/// ```
 #[derive(Clone, Copy, Debug)]
-pub struct TensorShape<const R: usize>([usize; R]);
+pub struct TensorShape<const R: usize>(pub(crate) [usize; R]);
+
+// ---- Construction and Initialization ---------------------------------------
 
 impl<const R: usize> TensorShape<R> {
+    /// Creates a new `TensorShape` instance.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3]);
+    /// assert_eq!(shape.dim_size(0), Some(&2));
+    /// assert_eq!(shape.dim_size(1), Some(&3));
+    /// assert_eq!(shape.dim_size(2), None);
+    /// ```
     pub const fn new(shape: [usize; R]) -> Self {
         Self(shape)
     }
+}
 
-    pub fn axis(&self, index: usize) -> Option<&usize> {
+// ---- Getters ---------------------------------------------------------------
+
+impl<const R: usize> TensorShape<R> {
+    /// Get the size of the specified dimension.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3]);
+    /// assert_eq!(shape.dim_size(0).unwrap(), &2);
+    /// assert_eq!(shape.dim_size(1).unwrap(), &3);
+    /// ```
+    pub fn dim_size(&self, index: usize) -> Option<&usize> {
         self.0.get(index)
     }
 
+    /// Returns the shape as an array.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3]);
+    /// assert_eq!(shape.as_array(), &[2, 3]);
+    /// ```
+    pub const fn as_array(&self) -> &[usize; R] {
+        &self.0
+    }
+
+    /// Returns the size of the shape, meaning the product of all dimensions.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3]);
+    /// assert_eq!(shape.size(), 6);
+    /// ```
+    pub fn size(&self) -> usize {
+        self.0.iter().product()
+    }
+}
+
+// ---- Manipulation ----------------------------------------------------------
+
+impl<const R: usize> TensorShape<R> {
+    /// Reorders the dimensions of the shape.
+    ///
+    /// ```
+    /// use manifold::*;
+    ///
+    /// let shape = TensorShape::new([2, 3, 4]);
+    /// let new_shape = shape.reorder([2, 0, 1]);
+    /// assert_eq!(new_shape, shape!([4, 2, 3]));
+    /// ```
     pub fn reorder(&self, indices: [usize; R]) -> Self {
         let mut new_shape = TensorShape::new([0; R]);
         for (new_index, &index) in indices.iter().enumerate() {
@@ -24,78 +96,15 @@ impl<const R: usize> TensorShape<R> {
         new_shape
     }
 
-    pub const fn as_array(&self) -> &[usize; R] {
+    /// Creates a new shape by removing the specified dimensions.
-        &self.0
-    }
-
-    pub const fn rank(&self) -> usize {
-        R
-    }
-
-    pub fn flat_max(&self) -> usize {
-        self.size() - 1
-    }
-
-    pub fn size(&self) -> usize {
-        self.0.iter().product()
-    }
-
-    pub fn iter(&self) -> impl Iterator<Item = &usize> {
-        self.0.iter()
-    }
-
-    pub const fn get(&self, index: usize) -> usize {
-        self.0[index]
-    }
-
-    pub fn check_indices(&self, indices: [usize; R]) -> bool {
-        indices
-            .iter()
-            .zip(self.0.iter())
-            .all(|(&idx, &dim_size)| idx < dim_size)
-    }
-
-    /// Converts a flat index to a multi-dimensional index.
     ///
-    /// # Arguments
+    /// ```
-    /// * `flat_index` - The flat index to convert.
+    /// use manifold::*;
     ///
-    /// # Returns
+    /// let shape = TensorShape::new([2, 3, 4]);
-    /// An `TensorIndex<R>` instance representing the multi-dimensional index
+    /// let new_shape = shape.remove_dims([0, 2]);
-    /// corresponding to the flat index.
+    /// assert_eq!(new_shape, shape!([3]));
-    ///
+    /// ```
-    /// # How It Works
-    /// - The method iterates over the dimensions of the tensor in reverse order
-    ///   (assuming row-major order).
-    /// - In each iteration, it uses the modulo operation to find the index in
-    ///   the current dimension and integer division to reduce the flat index
-    ///   for the next higher dimension.
-    /// - This process is repeated for each dimension to build the
-    ///   multi-dimensional index.
-    pub fn index_from_flat(&self, flat_index: usize) -> TensorIndex<R> {
-        let mut indices = [0; R];
-        let mut remaining = flat_index;
-
-        for (idx, &dim_size) in indices.iter_mut().zip(self.0.iter()).rev() {
-            *idx = remaining % dim_size;
-            remaining /= dim_size;
-        }
-
-        indices.reverse(); // Reverse the indices to match the original dimension order
-        TensorIndex::new(self.clone(), indices)
-    }
-
-    pub fn index_zero(&self) -> TensorIndex<R> {
-        TensorIndex::zero(self.clone())
-    }
-
-    pub fn index_max(&self) -> TensorIndex<R> {
-        let max_indices =
-            self.0
-                .map(|dim_size| if dim_size > 0 { dim_size - 1 } else { 0 });
-        TensorIndex::new(self.clone(), max_indices)
-    }
-
     pub fn remove_dims<const NAX: usize>(
         &self,
         dims_to_remove: [usize; NAX],
@@ -118,34 +127,70 @@ impl<const R: usize> TensorShape<R> {
 
         TensorShape(new_shape)
     }
+}
 
-    pub fn remove_axes<'a, T: Value, const NAX: usize>(
+// ---- Iterators -------------------------------------------------------------
-        &self,
-        axes_to_remove: &'a [TensorAxis<'a, T, R>; NAX],
-    ) -> TensorShape<{ R - NAX }> {
-        // Create a new array to store the remaining dimensions
-        let mut new_shape = [0; R - NAX];
-        let mut new_index = 0;
 
-        // Iterate over the original dimensions
+impl<const R: usize> TensorShape<R> {
-        for (index, &dim) in self.0.iter().enumerate() {
+    /// Returns an iterator over the dimensions of the shape.
-            // Skip dimensions that are in the axes_to_remove array
+    ///
-            for axis in axes_to_remove {
+    /// ```
-                if *axis.dim() == index {
+    /// use manifold::*;
-                    continue;
+    ///
+    /// let shape = TensorShape::new([2, 3]);
+    /// let mut iter = shape.iter();
+    /// assert_eq!(iter.next(), Some(&2));
+    /// assert_eq!(iter.next(), Some(&3));
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    pub fn iter(&self) -> impl Iterator<Item = &usize> {
+        self.0.iter()
     }
 }
 
-            // Add the dimension to the new shape array
+// ---- Utils -----------------------------------------------------------------
-            new_shape[new_index] = dim;
+
-            new_index += 1;
+impl<const R: usize> TensorShape<R> {
+    pub fn index_zero(&self) -> TensorIndex<R> {
+        TensorIndex::zero(self.clone())
     }
 
-        TensorShape(new_shape)
+    pub fn index_max(&self) -> TensorIndex<R> {
+        let max_indices =
+            self.0
+                .map(|dim_size| if dim_size > 0 { dim_size - 1 } else { 0 });
+        TensorIndex::new(self.clone(), max_indices)
     }
 }
 
-// ---- Serialize and Deserialize ----
+// ---- From ------------------------------------------------------------------
+
+impl<const R: usize> From<[usize; R]> for TensorShape<R> {
+    fn from(shape: [usize; R]) -> Self {
+        Self::new(shape)
+    }
+}
+
+impl<T, const R: usize> From<Tensor<T, R>> for TensorShape<R>
+where
+    T: Value,
+{
+    fn from(tensor: Tensor<T, R>) -> Self {
+        *tensor.shape()
+    }
+}
+
+// ---- Equality --------------------------------------------------------------
+
+impl<const R: usize> PartialEq for TensorShape<R> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0 == other.0
+    }
+}
+
+impl<const R: usize> Eq for TensorShape<R> {}
+
+// ---- Serialization ---------------------------------------------------------
 
 struct TensorShapeVisitor<const R: usize>;
 
@@ -191,30 +236,3 @@ impl<const R: usize> Serialize for TensorShape<R> {
         seq.end()
     }
 }
-
-// ---- Blanket Implementations ----
-
-impl<const R: usize> From<[usize; R]> for TensorShape<R> {
-    fn from(shape: [usize; R]) -> Self {
-        Self::new(shape)
-    }
-}
-
-impl<const R: usize> PartialEq for TensorShape<R> {
-    fn eq(&self, other: &Self) -> bool {
-        self.0 == other.0
-    }
-}
-
-impl<const R: usize> Eq for TensorShape<R> {}
-
-// ---- From and Into Implementations ----
-
-impl<T, const R: usize> From<Tensor<T, R>> for TensorShape<R>
-where
-    T: Value,
-{
-    fn from(tensor: Tensor<T, R>) -> Self {
-        *tensor.shape()
-    }
-}

src/tensor.rs

@@ -28,6 +28,9 @@ pub struct Tensor<T, const R: usize> {
     shape: TensorShape<R>,
 }
 
+/// A type that represents an axis of a tensor.
+pub struct Axis(pub usize);
+
 // ---- Construction and Initialization ---------------------------------------
 
 impl<T: Value, const R: usize> Tensor<T, R> {
@@ -62,22 +65,49 @@ impl<T: Value, const R: usize> Tensor<T, R> {
     /// use manifold::Tensor;
     ///
     /// let buffer = vec![1, 2, 3, 4, 5, 6];
-    /// let t = Tensor::<i32, 2>::new_with_buffer([2, 3].into(), buffer);
+    /// let t = Tensor::<i32, 2>::new_with_buffer([2, 3].into(), buffer).unwrap();
     /// assert_eq!(t.shape().as_array(), &[2, 3]);
     /// assert_eq!(t.buffer(), &[1, 2, 3, 4, 5, 6]);
     /// ```
-    pub fn new_with_buffer(shape: TensorShape<R>, buffer: Vec<T>) -> Self {
+    pub fn new_with_buffer(
-        Self { buffer, shape }
+        shape: TensorShape<R>,
+        buffer: Vec<T>,
+    ) -> Result<Self> {
+        if buffer.len() != shape.size() {
+            Err(TensorError::InvalidArgument(format!(
+                "Provided buffer has length {} but shape has size {}",
+                buffer.len(),
+                shape.size()
+            )))
+        } else {
+            Ok(Self { buffer, shape })
+        }
     }
 }
 
 // ---- Trivial Getters -------------------------------------------------------
 
 impl<T: Value, const R: usize> Tensor<T, R> {
-    pub fn rank(&self) -> usize {
+    /// Get the rank of the tensor.
+    ///
+    /// ```
+    /// use manifold::Tensor;
+    ///
+    /// let t = Tensor::<f64, 2>::new([3, 3].into());
+    /// assert_eq!(t.rank(), 2);
+    /// ```
+    pub const fn rank(&self) -> usize {
         R
     }
 
+    /// Get the length of the tensor's buffer.
+    ///
+    /// ```
+    /// use manifold::Tensor;
+    ///
+    /// let t = Tensor::<f64, 2>::new([3, 3].into());
+    /// assert_eq!(t.len(), 9);
+    /// ```
     pub fn len(&self) -> usize {
         self.buffer().len()
     }
@@ -332,32 +362,6 @@ impl<T: Value, const R: usize> Tensor<T, R> {
     ) -> Result<()> {
         self.ew_for_each(other, result, &|a, b| a % b)
     }
-
-    // pub fn product<const S: usize>(
-    //     &self,
-    //     other: &Tensor<T, S>,
-    // ) -> Tensor<T, { R + S }> {
-    //     let mut new_shape_vec = Vec::new();
-    //     new_shape_vec.extend_from_slice(&self.shape().as_array());
-    //     new_shape_vec.extend_from_slice(&other.shape().as_array());
-
-    //     let new_shape_array: [usize; R + S] = new_shape_vec
-    //         .try_into()
-    //         .expect("Failed to create shape array");
-
-    //     let mut new_buffer =
-    //         Vec::with_capacity(self.buffer.len() * other.buffer.len());
-    //     for &item_self in &self.buffer {
-    //         for &item_other in &other.buffer {
-    //             new_buffer.push(item_self * item_other);
-    //         }
-    //     }
-
-    //     Tensor {
-    //         buffer: new_buffer,
-    //         shape: TensorShape::new(new_shape_array),
-    //     }
-    // }
 }
 
 // ---- Reshape ---------------------------------------------------------------
@@ -385,37 +389,11 @@ impl<T: Value, const R: usize> Tensor<T, R> {
                 "TensorShape size mismatch: ( {ls:?} = {lsize} ) != ( {rs:?} = {rsize} )",
             )))
         } else {
-            Ok(Tensor::new_with_buffer(shape, self.buffer))
+            Ok(Tensor::new_with_buffer(shape, self.buffer).unwrap())
         }
     }
 }
 
-// ---- Transpose -------------------------------------------------------------
-
-impl<T: Value, const R: usize> Tensor<T, R> {
-    /// Transpose the tensor according to the given order. The order must be a
-    /// permutation of the tensor's axes.
-    ///
-    /// ```
-    /// use manifold::{tensor, Tensor, TensorShape};
-    ///
-    /// let t = tensor!([[1, 2], [3, 4]]);
-    /// let t = t.transpose([1, 0]).unwrap();
-    /// assert_eq!(t, tensor!([[1, 3], [2, 4]]));
-    /// ```
-    pub fn transpose(self, order: [usize; R]) -> Result<Self> {
-        let buffer = TensorIndex::from(self.shape().clone())
-            .iter_transposed(order)
-            .map(|index| self.get(index).unwrap().clone())
-            .collect();
-
-        Ok(Tensor {
-            buffer,
-            shape: self.shape().reorder(order),
-        })
-    }
-}
-
 // ---- Operations ------------------------------------------------------------
 
 impl<T: Value, const R: usize> Add for Tensor<T, R> {
@@ -654,73 +632,6 @@ where
 
 impl<T, const R: usize> Eq for Tensor<T, R> where T: Eq {}
 
-// ---- Iterator --------------------------------------------------------------
-
-pub struct TensorIterator<'a, T: Value, const R: usize> {
-    tensor: &'a Tensor<T, R>,
-    index: TensorIndex<R>,
-}
-
-impl<'a, T: Value, const R: usize> TensorIterator<'a, T, R> {
-    pub fn new(tensor: &'a Tensor<T, R>) -> Self {
-        Self {
-            tensor,
-            index: tensor.shape.index_zero(),
-        }
-    }
-}
-
-impl<'a, T: Value, const R: usize> Iterator for TensorIterator<'a, T, R> {
-    type Item = &'a T;
-    fn next(&mut self) -> Option<Self::Item> {
-        if self.index.is_overflow() {
-            return None;
-        }
-
-        let result = unsafe { self.tensor.get_unchecked(self.index) };
-        self.index.inc();
-        Some(result)
-    }
-}
-
-impl<'a, T: Value, const R: usize> IntoIterator for &'a Tensor<T, R> {
-    type Item = &'a T;
-    type IntoIter = TensorIterator<'a, T, R>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        TensorIterator::new(self)
-    }
-}
-
-// ---- Formatting ------------------------------------------------------------
-
-impl<'a, T: Value, const R: usize> Display for TensorIterator<'a, T, R> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
-        // Print the current index and flat index
-        write!(
-            f,
-            "Current Index: {}, Flat Index: {}",
-            self.index,
-            self.index.flat()
-        )?;
-
-        // Print the tensor elements, highlighting the current element
-        write!(f, ", Tensor Elements: [")?;
-        for (i, elem) in self.tensor.buffer().iter().enumerate() {
-            if i == self.index.flat() {
-                write!(f, "*{}*", elem)?; // Highlight the current element
-            } else {
-                write!(f, "{}", elem)?;
-            }
-            if i < self.tensor.buffer().len() - 1 {
-                write!(f, ", ")?;
-            }
-        }
-
-        write!(f, "]")
-    }
-}
-
 // ---- From ------------------------------------------------------------------
 
 impl<T: Value, const R: usize> From<TensorShape<R>> for Tensor<T, R> {

src/value.rs

@@ -1,22 +0,0 @@
-use num::{Num, One, Zero};
-use serde::{Deserialize, Serialize};
-use std::{fmt::Display, iter::Sum};
-
-/// A trait for types that can be used as values in a tensor.
-pub trait Value:
-    Num + Zero + One + Copy + Clone + Display + Serialize + Deserialize<'static>
-{
-}
-
-impl<T> Value for T where
-    T: Num
-        + Zero
-        + One
-        + Copy
-        + Clone
-        + Display
-        + Serialize
-        + Deserialize<'static>
-        + Sum
-{
-}

tests/basic_tests.rs

@@ -53,14 +53,10 @@ fn test_iterating_3d_tensor() {
         }
     }
 
-    println!("{}", tensor);
-
     // Iterate over the tensor and check that the numbers are correct
 
     let mut iter = TensorIterator::new(&tensor);
 
-    println!("{}", iter);
-
     assert_eq!(iter.next(), Some(&0));
 
     assert_eq!(iter.next(), Some(&1));
@@ -157,105 +153,105 @@ fn test_index_dec_method() {
     assert_eq!(index, TensorIndex::zero(shape));
 }
 
-#[test]
+// #[test]
-fn test_axis_iterator() {
+// fn test_axis_iterator() {
-    // Creating a 2x2 Tensor for testing
+//     // Creating a 2x2 Tensor for testing
-    let tensor = Tensor::from([[1.0, 2.0], [3.0, 4.0]]);
+//     let tensor = Tensor::from([[1.0, 2.0], [3.0, 4.0]]);
 
-    // Testing iteration over the first axis (axis = 0)
+//     // Testing iteration over the first axis (axis = 0)
-    let axis = TensorAxis::new(&tensor, 0);
+//     let axis = TensorAxis::new(&tensor, 0);
 
-    let mut axis_iter = axis.into_iter();
+//     let mut axis_iter = axis.into_iter();
 
-    assert_eq!(axis_iter.next(), Some(&1.0));
+//     assert_eq!(axis_iter.next(), Some(&1.0));
-    assert_eq!(axis_iter.next(), Some(&2.0));
+//     assert_eq!(axis_iter.next(), Some(&2.0));
-    assert_eq!(axis_iter.next(), Some(&3.0));
+//     assert_eq!(axis_iter.next(), Some(&3.0));
-    assert_eq!(axis_iter.next(), Some(&4.0));
+//     assert_eq!(axis_iter.next(), Some(&4.0));
 
-    // Resetting the iterator for the second axis (axis = 1)
+//     // Resetting the iterator for the second axis (axis = 1)
-    let axis = TensorAxis::new(&tensor, 1);
+//     let axis = TensorAxis::new(&tensor, 1);
 
-    let mut axis_iter = axis.into_iter();
+//     let mut axis_iter = axis.into_iter();
 
-    assert_eq!(axis_iter.next(), Some(&1.0));
+//     assert_eq!(axis_iter.next(), Some(&1.0));
-    assert_eq!(axis_iter.next(), Some(&3.0));
+//     assert_eq!(axis_iter.next(), Some(&3.0));
-    assert_eq!(axis_iter.next(), Some(&2.0));
+//     assert_eq!(axis_iter.next(), Some(&2.0));
-    assert_eq!(axis_iter.next(), Some(&4.0));
+//     assert_eq!(axis_iter.next(), Some(&4.0));
 
-    let shape = tensor.shape();
+//     let shape = tensor.shape();
 
-    let mut a: TensorIndex<2> = (shape.clone(), [0, 0]).into();
+//     let mut a: TensorIndex<2> = (shape.clone(), [0, 0]).into();
-    let b: TensorIndex<2> = (shape.clone(), [1, 1]).into();
+//     let b: TensorIndex<2> = (shape.clone(), [1, 1]).into();
 
-    while a <= b {
+//     while a <= b {
-        println!("a: {}", a);
+//         println!("a: {}", a);
-        a.inc();
+//         a.inc();
-    }
+//     }
-}
+// }
 
-#[test]
+// #[test]
-fn test_3d_tensor_axis_iteration() {
+// fn test_3d_tensor_axis_iteration() {
-    // Create a 3D Tensor with specific values
+//     // Create a 3D Tensor with specific values
-    // Tensor shape is 2x2x2 for simplicity
+//     // Tensor shape is 2x2x2 for simplicity
-    let t = Tensor::from([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]);
+//     let t = Tensor::from([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]);
 
-    // 	TensorAxis 0 (Layer-wise):
+//     // 	TensorAxis 0 (Layer-wise):
-    //
+//     //
-    // t[0][0][0] = 1
+//     // t[0][0][0] = 1
-    // t[0][0][1] = 2
+//     // t[0][0][1] = 2
-    // t[0][1][0] = 3
+//     // t[0][1][0] = 3
-    // t[0][1][1] = 4
+//     // t[0][1][1] = 4
-    // t[1][0][0] = 5
+//     // t[1][0][0] = 5
-    // t[1][0][1] = 6
+//     // t[1][0][1] = 6
-    // t[1][1][0] = 7
+//     // t[1][1][0] = 7
-    // t[1][1][1] = 8
+//     // t[1][1][1] = 8
-    // [1, 2, 3, 4, 5, 6, 7, 8]
+//     // [1, 2, 3, 4, 5, 6, 7, 8]
-    //
+//     //
-    // This order suggests that for each "layer" (first level of arrays),
+//     // This order suggests that for each "layer" (first level of arrays),
-    // the iterator goes through all rows and columns. It first completes
+//     // the iterator goes through all rows and columns. It first completes
-    // the entire first layer, then moves to the second.
+//     // the entire first layer, then moves to the second.
 
-    let a0 = TensorAxis::new(&t, 0);
+//     let a0 = TensorAxis::new(&t, 0);
-    let a0_order = a0.into_iter().cloned().collect::<Vec<_>>();
+//     let a0_order = a0.into_iter().cloned().collect::<Vec<_>>();
-    assert_eq!(a0_order, [1, 2, 3, 4, 5, 6, 7, 8]);
+//     assert_eq!(a0_order, [1, 2, 3, 4, 5, 6, 7, 8]);
 
-    // TensorAxis 1 (Row-wise within each layer):
+//     // TensorAxis 1 (Row-wise within each layer):
-    //
+//     //
-    // t[0][0][0] = 1
+//     // t[0][0][0] = 1
-    // t[0][0][1] = 2
+//     // t[0][0][1] = 2
-    // t[1][0][0] = 5
+//     // t[1][0][0] = 5
-    // t[1][0][1] = 6
+//     // t[1][0][1] = 6
-    // t[0][1][0] = 3
+//     // t[0][1][0] = 3
-    // t[0][1][1] = 4
+//     // t[0][1][1] = 4
-    // t[1][1][0] = 7
+//     // t[1][1][0] = 7
-    // t[1][1][1] = 8
+//     // t[1][1][1] = 8
-    // [1, 2, 5, 6, 3, 4, 7, 8]
+//     // [1, 2, 5, 6, 3, 4, 7, 8]
-    //
+//     //
-    // This indicates that within each "layer", the iterator first
+//     // This indicates that within each "layer", the iterator first
-    // completes the first row across all layers, then the second row
+//     // completes the first row across all layers, then the second row
-    // across all layers.
+//     // across all layers.
 
-    let a1 = TensorAxis::new(&t, 1);
+//     let a1 = TensorAxis::new(&t, 1);
-    let a1_order = a1.into_iter().cloned().collect::<Vec<_>>();
+//     let a1_order = a1.into_iter().cloned().collect::<Vec<_>>();
-    assert_eq!(a1_order, [1, 2, 5, 6, 3, 4, 7, 8]);
+//     assert_eq!(a1_order, [1, 2, 5, 6, 3, 4, 7, 8]);
 
-    // TensorAxis 2 (Column-wise within each layer):
+//     // TensorAxis 2 (Column-wise within each layer):
-    //
+//     //
-    // t[0][0][0] = 1
+//     // t[0][0][0] = 1
-    // t[0][1][0] = 3
+//     // t[0][1][0] = 3
-    // t[1][0][0] = 5
+//     // t[1][0][0] = 5
-    // t[1][1][0] = 7
+//     // t[1][1][0] = 7
-    // t[0][0][1] = 2
+//     // t[0][0][1] = 2
-    // t[0][1][1] = 4
+//     // t[0][1][1] = 4
-    // t[1][0][1] = 6
+//     // t[1][0][1] = 6
-    // t[1][1][1] = 8
+//     // t[1][1][1] = 8
-    // [1, 3, 5, 7, 2, 4, 6, 8]
+//     // [1, 3, 5, 7, 2, 4, 6, 8]
-    //
+//     //
-    // This indicates that within each "layer", the iterator first
+//     // This indicates that within each "layer", the iterator first
-    // completes the first column across all layers, then the second
+//     // completes the first column across all layers, then the second
-    // column across all layers.
+//     // column across all layers.
 
-    let a2 = TensorAxis::new(&t, 2);
+//     let a2 = TensorAxis::new(&t, 2);
-    let a2_order = a2.into_iter().cloned().collect::<Vec<_>>();
+//     let a2_order = a2.into_iter().cloned().collect::<Vec<_>>();
-    assert_eq!(a2_order, [1, 3, 5, 7, 2, 4, 6, 8]);
+//     assert_eq!(a2_order, [1, 3, 5, 7, 2, 4, 6, 8]);
-}
+// }