fix: fix some layout.dim_select(axis, 0) to layout.dim_chop(axis) to avoid panic on zero-sized tensors (#86)

Author: ajz34

rstsr-core/src/tensor/adv_indexing.rs

@@ -231,6 +231,13 @@ mod test {
             let b = a.index_select(0, [0, 0, 1, -1]);
             assert!(fingerprint(&b) - -25.648600916145096 < 1e-8);
         }
+
+        // 1-dim select with empty index
+        let device = DeviceCpuSerial::default();
+        let a = linspace((1.0, 4.0, 4, &device));
+        let mask: Vec<usize> = vec![];
+        let b = a.index_select(0, &mask);
+        assert_eq!(b.raw(), &[]);
     }
 
     #[test]

rstsr-core/src/tensor/pack_array.rs

@@ -145,7 +145,7 @@ where
         let (data, device) = storage.into_raw_parts();
         let data = data.pack_array_f::<N>()?;
         let storage = Storage::new(data, device);
-        let layout = layout.dim_select(axis as isize, 0)?;
+        let layout = layout.dim_chop(axis as isize)?;
         let stride = layout
             .stride()
             .as_ref()

rstsr-native-impl/src/cpu_rayon/adv_indexing.rs

@@ -31,8 +31,8 @@ where
     let size_indices = indices.len();
 
     if axis_contig_c {
-        let lc_rest = lc.clone().dim_select(axis as isize, 0)?;
-        let la_rest = la.clone().dim_select(axis as isize, 0)?;
+        let lc_rest = lc.clone().dim_chop(axis as isize)?;
+        let la_rest = la.clone().dim_chop(axis as isize)?;
         let layouts_rest = translate_to_col_major(&[&lc_rest, &la_rest], TensorIterOrder::K)?;
         let lc_rest = &layouts_rest[0];
         let la_rest = &layouts_rest[1];

rstsr-native-impl/src/cpu_serial/adv_indexing.rs

@@ -36,8 +36,8 @@ where
     let size_indices = indices.len();
 
     if axis_contig_c {
-        let lc_rest = lc.clone().dim_select(axis as isize, 0)?;
-        let la_rest = la.clone().dim_select(axis as isize, 0)?;
+        let lc_rest = lc.clone().dim_chop(axis as isize)?;
+        let la_rest = la.clone().dim_chop(axis as isize)?;
         let layouts_rest = translate_to_col_major(&[&lc_rest, &la_rest], TensorIterOrder::K)?;
         let lc_rest = &layouts_rest[0];
         let la_rest = &layouts_rest[1];