add permute operation to tensor

tensor/src/lib.rs

@@ -62,6 +62,14 @@ impl Tensor {
         Ok(())
     }
 
+    pub fn permute(&mut self, order: &[usize]) -> Result<(), &'static str> {
+        if order.len() != self.shape.len() {
+            return Err("The permutation does not align with the current shape.");
+        }
+        let new_shape: Vec<usize> = order.iter().map(|i| self.shape[*i]).collect();
+        self.reshape(&new_shape)
+    }
+
     pub fn shape(&self) -> &Vec<usize> {
         &self.shape
     }

tensor/tests/tensor_core_test.rs

@@ -1,3 +1,5 @@
+use std::usize;
+
 use tensor::Tensor;
 
 #[test]
@@ -64,6 +66,31 @@ fn reshape_tensor_invalid_shape() {
     }
 }
 
+#[test]
+fn permute_tensor_valid_order() {
+    let original_shape: Vec<usize> = vec![4, 2, 2];
+    let mut a: Tensor = Tensor::ones(&original_shape);
+
+    let permutation: Vec<usize> = vec![1, 0, 2];
+    let new_strides: Vec<usize> = vec![8, 2, 1];
+    let new_shape: Vec<usize> = vec![2, 4, 2];
+    a.permute(&permutation).unwrap();
+
+    assert_eq!(new_strides, *a.strides());
+    assert_eq!(new_shape, *a.shape());
+}
+
+#[test]
+fn permute_tensor_invalid_order() {
+    let original_shape: Vec<usize> = vec![4, 2];
+    let mut a: Tensor = Tensor::ones(&original_shape);
+
+    let permutation: Vec<usize> = vec![2, 0, 1];
+    if let Ok(()) = a.permute(&permutation) {
+        panic!("The permutation should've been invalid.")
+    }
+}
+
 #[test]
 fn get_element_with_index() {
     let length: usize = 24;