feat!: make dims private in Shape

Author: siliconlad

src/lib.rs

@@ -32,19 +32,12 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     }
 
     pub fn fill(shape: &Shape, value: T) -> Tensor<T> {
-        let total_size = shape.size();
-        let mut vec = Vec::with_capacity(total_size);
-        for _ in 0..total_size { vec.push(value); }
+        let mut vec = Vec::with_capacity(shape.size());
+        for _ in 0..shape.size() { vec.push(value); }
         Tensor::new(shape, &vec).unwrap()
     }
-
-    pub fn zeros(shape: &Shape) -> Tensor<T> {
-        Tensor::fill(shape, T::zero())
-    }
-
-    pub fn ones(shape: &Shape) -> Tensor<T> {
-        Tensor::fill(shape, T::one())
-    }
+    pub fn zeros(shape: &Shape) -> Tensor<T> {Tensor::fill(shape, T::zero())}
+    pub fn ones(shape: &Shape) -> Tensor<T> {Tensor::fill(shape, T::one())}
 
     // Properties
     pub fn shape(&self) -> &Shape { &self.shape }
@@ -64,8 +57,8 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     pub fn sum(&self, axes: Axes) -> Tensor<T> {
         let all_axes = (0..self.shape.order()).collect::<Vec<_>>();
         let remaining_axes = all_axes.clone().into_iter().filter(|&i| !axes.contains(&i)).collect::<Vec<_>>();
-        let remaining_dims = remaining_axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
-        let removing_dims = axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
+        let remaining_dims = remaining_axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
+        let removing_dims = axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
 
         // We resolve to a scalar value
         if axes.is_empty() | (remaining_dims.len() == 0) {
@@ -95,7 +88,7 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     }
 
     pub fn mean(&self, axes: Axes) -> Tensor<T> {
-        let removing_dims = axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
+        let removing_dims = axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
         let removing_dims_t: Vec<T> = removing_dims.iter().map(|&dim| {
             let mut result = T::zero();
             for _ in 0..dim {
@@ -108,7 +101,7 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     }
 
     pub fn var(&self, axes: Axes) -> Tensor<T> {
-        let removing_dims = axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
+        let removing_dims = axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
         let removing_dims_t: Vec<T> = removing_dims.iter().map(|&dim| {
             let mut result = T::zero();
             for _ in 0..dim {
@@ -120,8 +113,8 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
 
         let all_axes = (0..self.shape.order()).collect::<Vec<_>>();
         let remaining_axes = all_axes.clone().into_iter().filter(|&i| !axes.contains(&i)).collect::<Vec<_>>();
-        let remaining_dims = remaining_axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
-        let removing_dims = axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
+        let remaining_dims = remaining_axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
+        let removing_dims = axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
 
         // We resolve to a scalar value
         if axes.is_empty() | (remaining_dims.len() == 0) {
@@ -157,8 +150,8 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     pub fn max(&self, axes: Axes) -> Tensor<T> {
         let all_axes = (0..self.shape.order()).collect::<Vec<_>>();
         let remaining_axes = all_axes.clone().into_iter().filter(|&i| !axes.contains(&i)).collect::<Vec<_>>();
-        let remaining_dims = remaining_axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
-        let removing_dims = axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
+        let remaining_dims = remaining_axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
+        let removing_dims = axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
 
         // We resolve to a scalar value
         if axes.is_empty() | (remaining_dims.len() == 0) {
@@ -192,8 +185,8 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     pub fn min(&self, axes: Axes) -> Tensor<T> {
         let all_axes = (0..self.shape.order()).collect::<Vec<_>>();
         let remaining_axes = all_axes.clone().into_iter().filter(|&i| !axes.contains(&i)).collect::<Vec<_>>();
-        let remaining_dims = remaining_axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
-        let removing_dims = axes.iter().map(|&i| self.shape.dims[i]).collect::<Vec<_>>();
+        let remaining_dims = remaining_axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
+        let removing_dims = axes.iter().map(|&i| self.shape[i]).collect::<Vec<_>>();
 
         // We resolve to a scalar value
         if axes.is_empty() | (remaining_dims.len() == 0) {
@@ -227,9 +220,7 @@ impl<T: Num + PartialOrd + Copy> Tensor<T> {
     // Tensor Product
     // Consistent with numpy.tensordot(a, b, axis=0)
     pub fn prod(&self, other: &Tensor<T>) -> Tensor<T> {
-        let mut new_dims = self.shape.dims.clone();
-        new_dims.extend(&other.shape.dims);
-        let new_shape = Shape::new(new_dims).unwrap();
+        let new_shape = self.shape.stack(&other.shape);
 
         let mut new_data = Vec::with_capacity(self.size() * other.size());
         for &a in &self.data {

src/shape.rs

@@ -5,7 +5,7 @@ use crate::error::ShapeError;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct Shape {
-    pub dims: Vec<usize>,
+    dims: Vec<usize>,
 }
 
 impl Shape {
@@ -23,6 +23,12 @@ impl Shape {
     pub fn order(&self) -> usize {
         self.dims.len()
     }
+
+    pub fn stack(&self, rhs: &Shape) -> Shape {
+        let mut new_dims = self.dims.clone();
+        new_dims.extend(rhs.dims.iter());
+        Shape { dims: new_dims }
+    }
 }
 
 impl Index<usize> for Shape {
@@ -33,6 +39,14 @@ impl Index<usize> for Shape {
     }
 }
 
+impl Index<std::ops::RangeFrom<usize>> for Shape {
+    type Output = [usize];
+
+    fn index(&self, index: std::ops::RangeFrom<usize>) -> &Self::Output {
+        &self.dims[index]
+    }
+}
+
 impl fmt::Display for Shape {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         use itertools::Itertools;
@@ -83,4 +97,13 @@ mod tests {
         let shape = shape![2, 3, 4].unwrap();
         assert_eq!(shape.dims, vec![2, 3, 4]);
     }
+
+    
+    #[test]
+    fn test_shape_extend() {
+        let shape1 = Shape::new(vec![2, 3]).unwrap();
+        let shape2 = Shape::new(vec![4, 5]).unwrap();
+        let extended_shape = shape1.stack(&shape2);
+        assert_eq!(extended_shape.dims, vec![2, 3, 4, 5]);
+    }
 }

src/storage.rs

@@ -22,14 +22,14 @@ impl<T> DynamicStorage<T> {
         }
 
         for (i, &dim) in coord.iter().enumerate() {
-            if dim >= shape.dims[i] {
+            if dim >= shape[i] {
                 return Err(ShapeError::new(format!("out of bounds for dimension {}", i).as_str()));
             }
         }
 
         let mut index = 0;
         for k in 0..shape.order() {
-            let stride = shape.dims[k+1..].iter().product::<usize>();
+            let stride = shape[k+1..].iter().product::<usize>();
             index += coord[k] * stride;
         }
         Ok(index)