add flatten operation to tensor

PaytonWebber · Dec 18, 2024 · 60ce1bc · 60ce1bc
1 parent 4e2db76
commit 60ce1bc
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/tensor/src/lib.rs b/tensor/src/lib.rs
@@ -77,6 +77,11 @@ impl Tensor {
         self.reshape(&new_shape)
     }
 
+    pub fn flatten(&mut self) {
+        self.shape = vec![1];
+        self.strides = vec![1];
+    }
+
     pub fn shape(&self) -> &Vec<usize> {
         &self.shape
     }

diff --git a/tensor/tests/tensor_core_test.rs b/tensor/tests/tensor_core_test.rs
@@ -102,6 +102,18 @@ fn permute_tensor_invalid_order() {
     }
 }
 
+#[test]
+fn flatten_tensor() {
+    let length: usize = 42;
+    let expected_data: Vec<f32> = vec![1.0_f32; length];
+    let mut a: Tensor = Tensor::ones(&[7, 6]);
+
+    a.flatten();
+    assert_eq!(vec![1], *a.shape());
+    assert_eq!(vec![1], *a.strides());
+    assert_eq!(expected_data, *a.data());
+}
+
 #[test]
 fn get_element_with_index() {
     let length: usize = 24;