add check for index out of range permute

tensor/src/lib.rs

@@ -66,6 +66,13 @@ impl Tensor {
         if order.len() != self.shape.len() {
             return Err("The permutation does not align with the current shape.");
         }
+
+        let mut sorted_order: Vec<usize> = order.to_vec();
+        sorted_order.sort();
+        if sorted_order != (0..self.shape.len()).collect::<Vec<_>>() {
+            return Err("Index out of range for shape.");
+        }
+
         let new_shape: Vec<usize> = order.iter().map(|i| self.shape[*i]).collect();
         self.reshape(&new_shape)
     }

tensor/tests/tensor_core_test.rs

@@ -80,6 +80,17 @@ fn permute_tensor_valid_order() {
     assert_eq!(new_shape, *a.shape());
 }
 
+#[test]
+fn permute_tensor_index_out_of_range() {
+    let original_shape: Vec<usize> = vec![4, 2];
+    let mut a: Tensor = Tensor::ones(&original_shape);
+
+    let permutation: Vec<usize> = vec![4, 2, 1];
+    if let Ok(()) = a.permute(&permutation) {
+        panic!("The permutation should've been invalid.")
+    }
+}
+
 #[test]
 fn permute_tensor_invalid_order() {
     let original_shape: Vec<usize> = vec![4, 2];