Improve Split operator implementation, add TensorBase::{slice_axis, slice_axis_mut}

rten-tensor/src/errors.rs

@@ -56,6 +56,10 @@ pub enum SliceError {
         range_ndim: usize,
     },
 
+    /// The slice spec specified an axis that is equal to or greater than the
+    /// dimension count.
+    InvalidAxis { axis: usize },
+
     /// An index in the slice spec is out of bounds for the corresponding tensor
     /// dimension.
     InvalidIndex {
@@ -105,6 +109,7 @@ impl Display for SliceError {
                     range_ndim, ndim
                 )
             }
+            SliceError::InvalidAxis { axis } => write!(f, "slice axis {} is invalid", axis),
             SliceError::InvalidIndex { axis, index, size } => write!(
                 f,
                 "slice index {} is invalid for axis ({}) of size {}",

rten-tensor/src/layout.rs

@@ -970,8 +970,21 @@ pub trait MutLayout: Layout + Clone {
     /// Slice the layout along a given axis.
     ///
     /// Returns a tuple of `(offset_range, sliced_layout)`.
-    fn slice_axis(&self, axis: usize, range: Range<usize>) -> (Range<usize>, Self) {
-        assert!(range.end >= range.start);
+    fn slice_axis(
+        &self,
+        axis: usize,
+        range: Range<usize>,
+    ) -> Result<(Range<usize>, Self), SliceError> {
+        if axis >= self.ndim() {
+            return Err(SliceError::InvalidAxis { axis });
+        }
+        if range.end < range.start || range.end > self.size(axis) {
+            return Err(SliceError::InvalidRange {
+                axis,
+                range: range.into(),
+                size: self.size(axis),
+            });
+        }
 
         let mut sliced_layout = self.clone();
         sliced_layout.resize_dim(axis, range.len());
@@ -982,7 +995,7 @@ pub trait MutLayout: Layout + Clone {
             let end_offset = start_offset + sliced_layout.min_data_len();
             start_offset..end_offset
         };
-        (range, sliced_layout)
+        Ok((range, sliced_layout))
     }
 
     /// Return a layout with all size-one dimensions removed.
@@ -1782,13 +1795,49 @@ mod tests {
             } = case;
 
             let layout = DynLayout::from_shape(shape);
-            let (offset_range, sliced_layout) = layout.slice_axis(axis, range);
+            let (offset_range, sliced_layout) = layout.slice_axis(axis, range).unwrap();
             assert_eq!(sliced_layout.shape(), sliced_shape);
             assert_eq!(sliced_layout.strides(), layout.strides());
             assert_eq!(offset_range, offsets);
         })
     }
 
+    #[test]
+    fn test_slice_axis_invalid() {
+        #[derive(Debug)]
+        struct Case<'a> {
+            shape: &'a [usize],
+            axis: usize,
+            range: Range<usize>,
+            expected: SliceError,
+        }
+
+        let cases = [
+            Case {
+                shape: &[1, 2, 3],
+                axis: 4,
+                range: 0..1,
+                expected: SliceError::InvalidAxis { axis: 4 },
+            },
+            Case {
+                shape: &[1, 2, 3],
+                axis: 0,
+                range: 0..2,
+                expected: SliceError::InvalidRange {
+                    axis: 0,
+                    range: (0..2).into(),
+                    size: 1,
+                },
+            },
+        ];
+
+        cases.test_each(|case| {
+            let layout = DynLayout::from_shape(case.shape);
+            let result = layout.slice_axis(case.axis, case.range.clone());
+            assert_eq!(result, Err(case.expected.clone()));
+        })
+    }
+
     #[test]
     fn test_slice_invalid() {
         #[derive(Debug)]

rten-tensor/src/tensor.rs

@@ -321,6 +321,15 @@ pub trait AsView: Layout {
         self.view().slice(range)
     }
 
+    /// Slice this tensor along a given axis.
+    fn slice_axis(
+        &self,
+        axis: usize,
+        range: Range<usize>,
+    ) -> TensorBase<ViewData<Self::Elem>, Self::Layout> {
+        self.view().slice_axis(axis, range)
+    }
+
     /// A variant of [`slice`](Self::slice) that returns a result
     /// instead of panicking.
     #[allow(clippy::type_complexity)]
@@ -791,14 +800,13 @@ impl<S: StorageMut, L: MutLayout> TensorBase<S, L> {
     }
 
     /// Slice this tensor along a given axis.
-    fn slice_axis_mut(
+    pub fn slice_axis_mut(
         &mut self,
         axis: usize,
         range: Range<usize>,
     ) -> TensorBase<ViewMutData<S::Elem>, L> {
-        let (offset_range, sliced_layout) = self.layout.slice_axis(axis, range.clone());
+        let (offset_range, sliced_layout) = self.layout.slice_axis(axis, range.clone()).unwrap();
         debug_assert_eq!(sliced_layout.size(axis), range.len());
-
         TensorBase {
             data: self.data.slice_mut(offset_range),
             layout: sliced_layout,
@@ -1486,6 +1494,16 @@ impl<'a, T, L: Clone + MutLayout> TensorBase<ViewData<'a, T>, L> {
         self.try_slice(range).expect("slice failed")
     }
 
+    /// Slice this tensor along a given axis.
+    pub fn slice_axis(&self, axis: usize, range: Range<usize>) -> TensorBase<ViewData<'a, T>, L> {
+        let (offset_range, sliced_layout) = self.layout.slice_axis(axis, range.clone()).unwrap();
+        debug_assert_eq!(sliced_layout.size(axis), range.len());
+        TensorBase {
+            data: self.data.slice(offset_range),
+            layout: sliced_layout,
+        }
+    }
+
     /// A variant of [`slice`](Self::slice) that returns a result
     /// instead of panicking.
     #[allow(clippy::type_complexity)]
@@ -1513,7 +1531,7 @@ impl<'a, T, L: Clone + MutLayout> TensorBase<ViewData<'a, T>, L> {
         }
     }
 
-    /// Divide this tensor into two mutable views along a given axis.
+    /// Divide this tensor into two views along a given axis.
     ///
     /// Returns a `(left, right)` tuple of views, where the `left` view
     /// contains the slice from `[0, mid)` along `axis` and the `right`
@@ -3458,6 +3476,25 @@ mod tests {
         assert_eq!(row.data().unwrap(), &[1, 2, 3]);
     }
 
+    #[test]
+    fn test_slice_axis() {
+        let data = NdTensor::from([[1, 2, 3], [4, 5, 6]]);
+        let row = data.slice_axis(0, 0..1);
+        let col = data.slice_axis(1, 1..2);
+        assert_eq!(row, data.slice((0..1, ..)));
+        assert_eq!(col, data.slice((.., 1..2)));
+    }
+
+    #[test]
+    fn test_slice_axis_mut() {
+        let mut data = NdTensor::from([[1, 2, 3], [4, 5, 6]]);
+        let mut row = data.slice_axis_mut(0, 0..1);
+        row.fill(8);
+        let mut col = data.slice_axis_mut(1, 1..2);
+        col.fill(9);
+        assert_eq!(data, NdTensor::from([[8, 9, 8], [4, 9, 6]]));
+    }
+
     #[test]
     fn test_slice_mut() {
         // Slice static-rank array. The rank of the slice is inferred.

src/ops/split.rs

@@ -1,5 +1,5 @@
 use rten_tensor::prelude::*;
-use rten_tensor::{NdTensorView, SliceItem, Tensor, TensorView};
+use rten_tensor::{NdTensorView, Tensor, TensorView};
 
 use crate::ops::{
     map_input, resolve_axis, static_dims, Input, InputList, OpError, Operator, OutputList,
@@ -29,19 +29,9 @@ pub fn split<T: Copy>(
         .iter()
         .map(|&split_size| {
             let split_size = split_size as usize;
-            let slice_range: Vec<SliceItem> = (0..input.ndim())
-                .map(|dim| {
-                    if dim == axis {
-                        (split_start..split_start + split_size).into()
-                    } else {
-                        SliceItem::full_range()
-                    }
-                })
-                .collect();
-
+            let split_range = split_start..split_start + split_size;
             split_start += split_size;
-
-            input.slice(slice_range.as_slice()).to_tensor_in(pool)
+            input.slice_axis(axis, split_range).to_tensor_in(pool)
         })
         .collect();