Add step support to tensor slice operations (tracel-ai#3748)

* Add step support to tensor slice operations

Extends tensor slicing to support custom step values, including negative steps for reversed slices, via a new SliceInfo struct. Updates all relevant backend, API, and macro code to use SliceInfo for slicing, and adapts documentation and trait signatures to reflect the new functionality. Backends that do not support steps panic if step != 1. This enables advanced slicing patterns such as strided and reversed selection.

* Add support for slicing tensors with step in tch backend

Refactored tensor slicing logic to handle arbitrary step values, including negative steps, in the tch backend. Updated bool, int, and float tensor ops to use the new slice_with_steps method, removing previous step=1 restriction and panic.

* Add support for slicing tensors with steps

Introduces the `slice_with_steps` function to enable slicing tensors with arbitrary step sizes, including negative steps, in `base.rs`. Updates `bool_tensor.rs`, `int_tensor.rs`, and `tensor.rs` to use the new slicing logic via `SliceInfo`, allowing more flexible and efficient tensor slicing operations.

* Fix formatting

* Add step slicing support in burn-cubecl

* Fix format

* Refactor tensor slicing to use SliceInfo

Replaces usage of std::ops::Range with burn_tensor::SliceInfo for tensor slicing operations in qtensor and ring modules. This change improves consistency and prepares for more flexible slicing semantics across backends.

* Fix format

* Fix formatting

* Refactor tensor slice ops to use SliceInfo

Updated bool, float, int, and quantized tensor slice operations to accept and process SliceInfo instead of Range<usize>. This change enables support for slicing with steps and improves consistency across tensor types.

* Refactor slice kernels to use LinearView

Updated slice and slice_with_steps kernels to operate on LinearView types instead of raw Tensor references. This change improves memory access safety and consistency, and updates kernel launch calls to pass linear views for both input and output tensors.

* Refactor tensor slice output shape calculation

Centralizes and simplifies output shape calculation for tensor slicing with steps by introducing `calculate_slice_output_shape` in burn-tensor. Updates all relevant tensor ops to use this utility, improving code maintainability and consistency. Adds comprehensive unit tests for the new function.

* Fix no-std

* Add support for non-unit steps in ONNX Slice op

This commit adds support for arbitrary step values (including negative steps) in the ONNX Slice operator, updating the parser, code generation, and tests. It introduces new test cases and models for slicing with steps, removes previous restrictions on steps, and ensures correct handling of steps in both tensor and shape slicing scenarios.

* Fix lint warning

* Refactor cubecl slice kernel

* Fix no-std

* Refactor tensor slice validation to use SliceInfo

* Refactor slice_dim to use SliceInfo parameter

Updated the slice_dim function to accept a SliceInfo parameter instead of a Range, allowing for more flexible slicing with custom step values. Adjusted documentation and internal usage to reflect this change.

* Enforce unit step in slice_assign and slice_fill ops

Updated tensor slicing APIs to require step=1 for slice_assign and slice_fill, panicking if non-unit steps are provided. Refactored internal usage to pass SliceInfo instead of Range, updated related ops (cat, repeat_dim), and added tests to verify panics on unsupported stepped slicing.

* Add ignored test for matmul with stepped slice

Introduces a new test 'should_diff_matmul_with_slice_stepped' to verify autodiff behavior when using stepped slices in matmul operations. The test is currently ignored due to lack of support for slice assignment with steps.

* Refactor tensor slicing to use Slice instead of SliceInfo

Replaces usage of SliceInfo with the new Slice type across all tensor slicing APIs, implementations, and internal logic. Updates function signatures, backend implementations, and related code to support the new Slice structure, improving consistency and flexibility for tensor slicing operations.

* Add panic for unsupported slice step in Autodiff

Added a check in float_slice to panic if any slice step is not 1, as Autodiff does not support slicing with step != 1. Also added a test to verify that slicing with a step triggers the expected panic.

* Refactor slice range checks and update tests

Improved slice range validation in TensorCheck by checking the raw end index before conversion, providing clearer error messages. Updated related tests to use the Slice API directly and replaced Slice::new with Slice::from_range_stepped for consistency.

* Rename slice_infos to slices in tensor APIs

Updated documentation and variable names from 'slice_infos' to 'slices' across tensor API traits and convolution module implementations for consistency and clarity. No functional changes were made.

* Update tensor slicing docs to use 'slices' terminology

* Refactor tensor slicing to use Slice::from with ranges

Replaces multiple usages of Slice::new with Slice::from and Rust range syntax for tensor slicing in convolution weight gradient functions. This improves code readability and leverages more idiomatic Rust constructs.

---------

Co-authored-by: Genna Wingert <[email protected]>

Author: antimora

burn-book/src/building-blocks/tensor.md

@@ -166,7 +166,7 @@ Those operations are available for all tensor kinds: `Int`, `Float`, and `Bool`.
 | `tensor.select(dim, indices)`               | `tensor.index_select(dim, indices)`                                       |
 | `tensor.select_assign(dim, indices, values)`| N/A                                                                       |
 | `tensor.shape()`                            | `tensor.shape`                                                            |
-| `tensor.slice(ranges)`                      | `tensor[(*ranges,)]`                                                      |
+| `tensor.slice(s![range;step])`              | `tensor[(*ranges,)]` or `tensor[start:end:step]`                          |
 | `tensor.slice_assign(ranges, values)`       | `tensor[(*ranges,)] = values`                                             |
 | `tensor.slice_fill(ranges, value)`          | `tensor[(*ranges,)] = value`                                              |
 | `tensor.slice_dim(dim, range)`              | N/A                                                                       |
@@ -181,7 +181,6 @@ Those operations are available for all tensor kinds: `Int`, `Float`, and `Bool`.
 | `tensor.unsqueeze_dim(dim)`                 | `tensor.unsqueeze(dim)`                                                   |
 | `tensor.unsqueeze_dims(dims)`               | N/A                                                                       |
 
-
 ### Numeric Operations
 
 Those operations are available for numeric tensor kinds: `Float` and `Int`.

crates/burn-autodiff/src/ops/bool_tensor.rs

@@ -32,8 +32,8 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
         B::bool_reshape(tensor, shape)
     }
 
-    fn bool_slice(tensor: BoolTensor<B>, ranges: &[core::ops::Range<usize>]) -> BoolTensor<B> {
-        B::bool_slice(tensor, ranges)
+    fn bool_slice(tensor: BoolTensor<B>, slices: &[burn_tensor::Slice]) -> BoolTensor<B> {
+        B::bool_slice(tensor, slices)
     }
 
     fn bool_empty(shape: Shape, device: &Device<B>) -> BoolTensor<B> {

crates/burn-autodiff/src/ops/int_tensor.rs

@@ -28,8 +28,8 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_reshape(tensor, shape)
     }
 
-    fn int_slice(tensor: IntTensor<B>, ranges: &[core::ops::Range<usize>]) -> IntTensor<B> {
-        B::int_slice(tensor, ranges)
+    fn int_slice(tensor: IntTensor<B>, slices: &[burn_tensor::Slice]) -> IntTensor<B> {
+        B::int_slice(tensor, slices)
     }
 
     fn int_empty(

crates/burn-autodiff/src/ops/qtensor.rs

@@ -1,5 +1,3 @@
-use core::ops::Range;
-
 use burn_tensor::{
     Device, Shape, TensorData,
     backend::Backend,
@@ -84,7 +82,10 @@ impl<B: Backend, C: CheckpointStrategy> QTensorOps<Self> for Autodiff<B, C> {
         unimplemented!()
     }
 
-    fn q_slice(_tensor: QuantizedTensor<Self>, _ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
+    fn q_slice(
+        _tensor: QuantizedTensor<Self>,
+        _slices: &[burn_tensor::Slice],
+    ) -> QuantizedTensor<Self> {
         unimplemented!()
     }
 

crates/burn-autodiff/src/ops/tensor.rs

@@ -1130,24 +1130,31 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
         }
     }
 
-    fn float_slice(
-        tensor: FloatTensor<Self>,
-        ranges: &[core::ops::Range<usize>],
-    ) -> FloatTensor<Self> {
+    fn float_slice(tensor: FloatTensor<Self>, slices: &[burn_tensor::Slice]) -> FloatTensor<Self> {
+        // Check if any slice has step != 1
+        for (i, slice) in slices.iter().enumerate() {
+            if slice.step != 1 {
+                panic!(
+                    "Autodiff does not support slice with step != 1. Dimension {} has step {}",
+                    i, slice.step
+                );
+            }
+        }
+
         #[derive(Debug)]
         struct Index;
 
         #[derive(new, Debug)]
         struct RetroSlice<B: Backend> {
             tensor_id: NodeID,
-            ranges: Vec<core::ops::Range<usize>>,
+            slices: Vec<burn_tensor::Slice>,
             _backend: PhantomData<B>,
         }
 
         impl<B: Backend> RetroForward for RetroSlice<B> {
             fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
                 let tensor = states.get_state::<B::FloatTensorPrimitive>(&self.tensor_id);
-                let out = B::float_slice(tensor, &self.ranges);
+                let out = B::float_slice(tensor, &self.slices);
                 states.save(out_node, out)
             }
         }
@@ -1170,22 +1177,30 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
             }
         }
 
+        // Convert slices to ranges for backward compatibility in State
+        let shape = tensor.primitive.shape();
+        let ranges: Vec<core::ops::Range<usize>> = slices
+            .iter()
+            .enumerate()
+            .map(|(i, s)| s.to_range(shape.dims[i]))
+            .collect();
+
         match Index
             .prepare::<C>([tensor.node.clone()])
             .memory_bound()
-            .retro_forward(RetroSlice::<B>::new(tensor.node.id, ranges.to_vec()))
+            .retro_forward(RetroSlice::<B>::new(tensor.node.id, slices.to_vec()))
             .parents([&tensor])
             .stateful()
         {
             OpsKind::Tracked(prep) => prep.finish(
                 (
-                    ranges.to_vec(),
+                    ranges,
                     tensor.primitive.shape(),
                     B::float_device(&tensor.primitive),
                 ),
-                B::float_slice(tensor.primitive, ranges),
+                B::float_slice(tensor.primitive, slices),
             ),
-            OpsKind::UnTracked(prep) => prep.finish(B::float_slice(tensor.primitive, ranges)),
+            OpsKind::UnTracked(prep) => prep.finish(B::float_slice(tensor.primitive, slices)),
         }
     }
 
@@ -1234,7 +1249,16 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
                         let zeros = B::float_zeros(shape_rhs, &device, grad.dtype().into());
                         B::float_slice_assign(grad, &ranges_4lhs.unwrap(), zeros)
                     },
-                    |grad| B::float_slice(grad, &ranges_4rhs.unwrap()),
+                    |grad| {
+                        let slices: Vec<burn_tensor::Slice> = ranges_4rhs
+                            .unwrap()
+                            .iter()
+                            .map(|r| {
+                                burn_tensor::Slice::new(r.start as isize, Some(r.end as isize), 1)
+                            })
+                            .collect();
+                        B::float_slice(grad, &slices)
+                    },
                 );
             }
         }
@@ -2122,7 +2146,13 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
                         let mut ranges = ranges.clone();
                         ranges[self.dim] = current_index..dim_size + current_index;
                         current_index += dim_size;
-                        grads.register::<B>(node.id, B::float_slice(grad.clone(), &ranges));
+                        let slices: Vec<burn_tensor::Slice> = ranges
+                            .iter()
+                            .map(|r| {
+                                burn_tensor::Slice::new(r.start as isize, Some(r.end as isize), 1)
+                            })
+                            .collect();
+                        grads.register::<B>(node.id, B::float_slice(grad.clone(), &slices));
                     });
             }
 

crates/burn-autodiff/src/tests/slice.rs

@@ -130,4 +130,46 @@ mod tests {
             .to_data()
             .assert_approx_eq::<FT>(&cat_grad_2.to_data(), Tolerance::default());
     }
+
+    #[test]
+    #[ignore = "slice_assign with steps are not supported currently"]
+    fn should_diff_matmul_with_slice_stepped() {
+        use burn_tensor::s;
+        let data_1 = TensorData::from([[1.0, 7.0], [100.0, 100.0], [2.0, 3.0], [100.0, 100.0]]);
+        let data_2 = TensorData::from([[4.0, 100.0, 7.0, 100.0], [2.0, 100.0, 3.0, 15.0]]);
+
+        let device = Default::default();
+        let tensor_1 = TestAutodiffTensor::<2>::from_data(data_1, &device).require_grad();
+        let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
+
+        let tensor_3 = tensor_1.clone().slice(s![0..;2, 0..2]); // [[1., 7.], [2., 3.]]
+        let tensor_4 = tensor_2.clone().slice(s![0..2, 0..;2]); // [[4., 7.], [2., 3.]]
+        let tensor_5 = tensor_3.clone().matmul(tensor_4);
+        let grads = tensor_5.backward();
+
+        let grad_1 = tensor_1.grad(&grads).unwrap();
+        let grad_2 = tensor_2.grad(&grads).unwrap();
+
+        grad_1.to_data().assert_eq(
+            &TensorData::from([[11., 5.], [0., 0.], [11., 5.], [0., 0.]]),
+            false,
+        );
+        grad_2.to_data().assert_eq(
+            &TensorData::from([[3., 0., 3., 0.], [10., 0., 10., 0.]]),
+            false,
+        );
+    }
+
+    #[test]
+    #[should_panic(expected = "Autodiff does not support slice with step != 1")]
+    fn should_panic_on_slice_with_step() {
+        use burn_tensor::s;
+
+        let data = TensorData::from([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]]);
+        let device = Default::default();
+        let tensor = TestAutodiffTensor::<2>::from_data(data, &device).require_grad();
+
+        // This should panic because step is 2
+        let _sliced = tensor.slice(s![.., 0..4; 2]);
+    }
 }

crates/burn-candle/src/ops/base.rs

@@ -105,6 +105,66 @@ pub fn slice(tensor: CandleTensor, ranges: &[std::ops::Range<usize>]) -> CandleT
     CandleTensor::new(narrow_tensor)
 }
 
+pub fn slice_with_steps(tensor: CandleTensor, slices: &[burn_tensor::Slice]) -> CandleTensor {
+    let mut result_tensor = tensor.tensor;
+
+    for (dim, slice) in slices.iter().enumerate() {
+        if slice.step == 1 {
+            // Use narrow for step=1 (more efficient)
+            // Convert slice to range using tensor shape
+            let dim_size = result_tensor.dim(dim).unwrap();
+            let range = slice.to_range(dim_size);
+            let start = range.start;
+            let length = range.end - range.start;
+            result_tensor = result_tensor.narrow(dim, start, length).unwrap();
+        } else {
+            // Use index_select for step != 1
+            let dim_size = result_tensor.dim(dim).unwrap();
+            let range = slice.to_range(dim_size);
+            let start = range.start;
+            let end = range.end;
+            let step = slice.step;
+
+            // Generate indices based on step direction
+            let indices_vec = if step > 0 {
+                // Forward stepping
+                let step_usize = step as usize;
+                (start..end).step_by(step_usize).collect::<Vec<_>>()
+            } else {
+                // Backward stepping (negative step)
+                let step_usize = step.unsigned_abs();
+                // Start from end-1 and go backwards
+                let mut indices = Vec::new();
+                let mut idx = end - 1;
+                while idx >= start && idx < end {
+                    // Check for underflow
+                    indices.push(idx);
+                    if idx >= step_usize {
+                        idx -= step_usize;
+                    } else {
+                        break;
+                    }
+                }
+                indices
+            };
+
+            // Convert indices to tensor and use index_select
+            let indices_len = indices_vec.len();
+            let device = result_tensor.device();
+            let indices = candle_core::Tensor::from_vec(
+                indices_vec.iter().map(|&x| x as u32).collect::<Vec<_>>(),
+                indices_len,
+                device,
+            )
+            .unwrap();
+
+            result_tensor = result_tensor.index_select(&indices, dim).unwrap();
+        }
+    }
+
+    CandleTensor::new(result_tensor)
+}
+
 pub fn slice_assign(
     tensor: CandleTensor,
     ranges: &[std::ops::Range<usize>],

crates/burn-candle/src/ops/bool_tensor.rs

@@ -57,8 +57,8 @@ impl<F: FloatCandleElement, I: IntCandleElement> BoolTensorOps<Self> for Candle<
         super::base::reshape(tensor, shape)
     }
 
-    fn bool_slice(tensor: BoolTensor<Self>, ranges: &[std::ops::Range<usize>]) -> BoolTensor<Self> {
-        super::base::slice(tensor, ranges)
+    fn bool_slice(tensor: BoolTensor<Self>, slices: &[burn_tensor::Slice]) -> BoolTensor<Self> {
+        super::base::slice_with_steps(tensor, slices)
     }
 
     fn bool_slice_assign(

crates/burn-candle/src/ops/int_tensor.rs

@@ -41,8 +41,8 @@ impl<F: FloatCandleElement, I: IntCandleElement> IntTensorOps<Self> for Candle<F
         super::base::reshape(tensor, shape)
     }
 
-    fn int_slice(tensor: IntTensor<Self>, indices: &[std::ops::Range<usize>]) -> IntTensor<Self> {
-        super::base::slice(tensor, indices)
+    fn int_slice(tensor: IntTensor<Self>, slices: &[burn_tensor::Slice]) -> IntTensor<Self> {
+        super::base::slice_with_steps(tensor, slices)
     }
 
     fn int_slice_assign(

crates/burn-candle/src/ops/qtensor.rs

@@ -1,5 +1,3 @@
-use std::ops::Range;
-
 use burn_tensor::{
     DType, Device, Shape, TensorData,
     backend::Backend,
@@ -80,7 +78,10 @@ impl<F: FloatCandleElement, I: IntCandleElement> QTensorOps<Self> for Candle<F,
         unimplemented!()
     }
 
-    fn q_slice(_tensor: QuantizedTensor<Self>, _ranges: &[Range<usize>]) -> QuantizedTensor<Self> {
+    fn q_slice(
+        _tensor: QuantizedTensor<Self>,
+        _slices: &[burn_tensor::Slice],
+    ) -> QuantizedTensor<Self> {
         unimplemented!()
     }