fix issue where flatten did not update shape correctly

PaytonWebber · Dec 18, 2024 · 210c13d · 210c13d
1 parent 60ce1bc
commit 210c13d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tensor/src/lib.rs b/tensor/src/lib.rs
@@ -78,7 +78,7 @@ impl Tensor {
     }
 
     pub fn flatten(&mut self) {
-        self.shape = vec![1];
+        self.shape = vec![self.shape.iter().product()];
         self.strides = vec![1];
     }
 

diff --git a/tensor/tests/tensor_core_test.rs b/tensor/tests/tensor_core_test.rs
@@ -109,9 +109,11 @@ fn flatten_tensor() {
     let mut a: Tensor = Tensor::ones(&[7, 6]);
 
     a.flatten();
-    assert_eq!(vec![1], *a.shape());
+    let elem: f32 = a[&[22]];
+    assert_eq!(vec![length], *a.shape());
     assert_eq!(vec![1], *a.strides());
     assert_eq!(expected_data, *a.data());
+    assert_eq!(elem, 1.0_f32);
 }
 
 #[test]