Fix repeat for dims > 1 (tracel-ai#1713)

crates/burn-fusion/src/ops/boolean.rs

@@ -564,7 +564,7 @@ impl<B: FusionBackend> BoolTensorOps<Self> for Fusion<B> {
 
  let stream = tensor.stream;
  let mut shape = tensor.shape.clone();
- shape[dim] = times;
+ shape[dim] *= times;
  let out = tensor.client.tensor_uninitialized(shape);
 
  let desc = RepeatOperationDescription {

crates/burn-fusion/src/ops/float.rs

@@ -1620,7 +1620,7 @@ impl<B: FusionBackend> FloatTensorOps<Self> for Fusion<B> {
 
  let stream = tensor.stream;
  let mut shape = tensor.shape.clone();
- shape[dim] = times;
+ shape[dim] *= times;
  let out = tensor.client.tensor_uninitialized(shape);
 
  let desc = RepeatOperationDescription {

crates/burn-fusion/src/ops/int.rs

@@ -1665,7 +1665,7 @@ impl<B: FusionBackend> IntTensorOps<Self> for Fusion<B> {
 
  let stream = tensor.stream;
  let mut shape = tensor.shape.clone();
- shape[dim] = times;
+ shape[dim] *= times;
  let out = tensor.client.tensor_uninitialized(shape);
 
  let desc = RepeatOperationDescription {

crates/burn-jit/src/kernel/index/repeat.rs

@@ -38,19 +38,21 @@ impl RepeatComputeShader {
 
  let stride_input = scope.create_local(Elem::UInt);
  let stride_output = scope.create_local(Elem::UInt);
- let shape_output = scope.create_local(Elem::UInt);
+ let shape = scope.create_local(Elem::UInt);
 
  for i in 0..self.rank {
+ gpu!(scope, stride_input = stride(input, i));
+ gpu!(scope, stride_output = stride(output, i));
  if i != self.dim {
- gpu!(scope, stride_input = stride(input, i));
- gpu!(scope, stride_output = stride(output, i));
- gpu!(scope, shape_output = shape(output, i));
-
- gpu!(scope, offset_local = id / stride_output);
- gpu!(scope, offset_local = offset_local % shape_output);
- gpu!(scope, offset_local = offset_local * stride_input);
- gpu!(scope, offset_input += offset_local);
+ gpu!(scope, shape = shape(output, i));
+ } else {
+ gpu!(scope, shape = shape(input, i));
  }
+
+ gpu!(scope, offset_local = id / stride_output);
+ gpu!(scope, offset_local = offset_local % shape);
+ gpu!(scope, offset_local = offset_local * stride_input);
+ gpu!(scope, offset_input += offset_local);
  }
 
  let result = scope.create_local(input.item());
@@ -108,12 +110,9 @@ pub(crate) fn repeat<R: Runtime, E: JitElement, const D1: usize>(
  times: usize,
 ) -> JitTensor<R, E, D1> {
  let mut shape = input.shape.clone();
- if shape.dims[dim] != 1 {
- panic!("Can only repeat dimension with dim=1");
- }
 
  // Create output handle
- shape.dims[dim] = times;
+ shape.dims[dim] *= times;
  let num_elems_output = shape.num_elements();
  let handle = input
  .client

crates/burn-tensor/src/tensor/api/base.rs

@@ -564,10 +564,6 @@ where
  }
 
  /// Repeat the tensor along the given dimension.
- ///
- /// # Panics
- ///
- /// If the selected dimension more than one item.
  pub fn repeat(self, dim: usize, times: usize) -> Self {
  Self::new(K::repeat(self.primitive, dim, times))
  }

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

@@ -1,4 +1,7 @@
-use super::{cat::cat_with_slice_assign, BoolTensor, Device, FloatTensor, IntTensor};
+use super::{
+ cat::cat_with_slice_assign, repeat::repeat_with_slice_assign, BoolTensor, Device, FloatTensor,
+ IntTensor,
+};
 use crate::{
  backend::Backend, chunk, narrow, tensor::Shape, Bool, Data, ElementConversion, Tensor,
 };
@@ -174,28 +177,12 @@ pub trait BoolTensorOps<B: Backend> {
  dim: usize,
  times: usize,
  ) -> BoolTensor<B, D> {
- let mut shape = Self::bool_shape(&tensor);
- if shape.dims[dim] != 1 {
- panic!("Can only repeat dimension with dim=1");
- }
- shape.dims[dim] = times;
-
- let mut i = 0;
- let ranges_select_all = [0; D].map(|_| {
- let start = 0;
- let end = shape.dims[i];
- i += 1;
- start..end
- });
-
- let mut tensor_output = Self::bool_empty(shape, &Self::bool_device(&tensor));
- for i in 0..times {
- let mut ranges = ranges_select_all.clone();
- ranges[dim] = i..i + 1;
- tensor_output = Self::bool_slice_assign(tensor_output, ranges, tensor.clone());
- }
-
- tensor_output
+ repeat_with_slice_assign::<B, D, Bool>(
+ Tensor::<B, D, Bool>::from_primitive(tensor),
+ dim,
+ times,
+ )
+ .into_primitive()
  }
 
  /// Concatenates the tensors along the given dimension.

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

@@ -1,4 +1,5 @@
 use super::cat::cat_with_slice_assign;
+use super::repeat::repeat_with_slice_assign;
 use super::{BoolTensor, Device, FloatTensor, IntElem, IntTensor};
 use crate::Tensor;
 use crate::{backend::Backend, tensor::Shape, Data, Distribution, ElementConversion, Int};
@@ -270,28 +271,12 @@ pub trait IntTensorOps<B: Backend> {
  dim: usize,
  times: usize,
  ) -> IntTensor<B, D> {
- let mut shape = Self::int_shape(&tensor);
- if shape.dims[dim] != 1 {
- panic!("Can only repeat dimension with dim=1");
- }
- shape.dims[dim] = times;
-
- let mut i = 0;
- let indices_select_all = [0; D].map(|_| {
- let start = 0;
- let end = shape.dims[i];
- i += 1;
- start..end
- });
-
- let mut tensor_output = Self::int_empty(shape, &Self::int_device(&tensor));
- for i in 0..times {
- let mut indices = indices_select_all.clone();
- indices[dim] = i..i + 1;
- tensor_output = Self::int_slice_assign(tensor_output, indices, tensor.clone());
- }
-
- tensor_output
+ repeat_with_slice_assign::<B, D, Int>(
+ Tensor::<B, D, Int>::from_primitive(tensor),
+ dim,
+ times,
+ )
+ .into_primitive()
  }
 
  /// Concatenates the given tensors along the given dimension.

crates/burn-tensor/src/tensor/ops/modules/cat.rs

@@ -19,10 +19,8 @@ pub(crate) fn cat_with_slice_assign<B: Backend, const D: usize, K: TensorKind<B>
 
  let mut i = 0;
  let indices_select_all = [0; D].map(|_| {
- let start = 0;
- let end = shape.dims[i];
  i += 1;
- start..end
+ 0..shape.dims[i - 1]
  });
 
  let mut output_index = 0;

crates/burn-tensor/src/tensor/ops/modules/mod.rs

@@ -3,6 +3,8 @@ pub mod conv;
 
 /// Module with cat operation
 pub(crate) mod cat;
+/// Module with repeat operation
+pub(crate) mod repeat;
 /// Module with unfold operations.
 pub(crate) mod unfold;
 

crates/burn-tensor/src/tensor/ops/modules/repeat.rs

@@ -0,0 +1,36 @@
+use crate::{backend::Backend, BasicOps, Tensor, TensorKind};
+
+pub(crate) fn repeat_with_slice_assign<
+ B: Backend,
+ const D: usize,
+ K: TensorKind<B> + BasicOps<B>,
+>(
+ tensor: Tensor<B, D, K>,
+ dim: usize,
+ times: usize,
+) -> Tensor<B, D, K> {
+ let mut shape = tensor.shape();
+ let device = tensor.device();
+
+ let original_dim_length = shape.dims[dim];
+ shape.dims[dim] *= times;
+
+ let mut tensor_output = Tensor::empty(shape.clone(), &device);
+
+ let mut i = 0;
+ let indices_select_all = [0; D].map(|_| {
+ i += 1;
+ 0..shape.dims[i - 1]
+ });
+
+ let mut output_index = 0;
+ for _ in 0..times {
+ let mut indices = indices_select_all.clone();
+ indices[dim] = output_index..output_index + original_dim_length;
+ output_index += original_dim_length;
+
+ tensor_output = tensor_output.slice_assign(indices, tensor.clone());
+ }
+
+ tensor_output
+}

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

@@ -1,4 +1,5 @@
 use super::cat::cat_with_slice_assign;
+use super::repeat::repeat_with_slice_assign;
 use super::{BoolTensor, Device, FloatElem, FloatTensor, FullPrecisionBackend, IntElem, IntTensor};
 use crate::backend::BackendBridge;
 use crate::Tensor;
@@ -193,28 +194,8 @@ pub trait FloatTensorOps<B: Backend> {
  dim: usize,
  times: usize,
  ) -> FloatTensor<B, D> {
- let mut shape = B::float_shape(&tensor);
- if shape.dims[dim] != 1 {
- panic!("Can only repeat dimension with dim=1");
- }
- shape.dims[dim] = times;
-
- let mut i = 0;
- let indices_select_all = [0; D].map(|_| {
- let start = 0;
- let end = shape.dims[i];
- i += 1;
- start..end
- });
-
- let mut tensor_output = B::float_empty(shape, &B::float_device(&tensor));
- for i in 0..times {
- let mut indices = indices_select_all.clone();
- indices[dim] = i..i + 1;
- tensor_output = B::float_slice_assign(tensor_output, indices, tensor.clone());
- }
-
- tensor_output
+ repeat_with_slice_assign::<B, D, Float>(Tensor::<B, D>::from_primitive(tensor), dim, times)
+ .into_primitive()
  }
 
  /// Adds two tensors together.

crates/burn-tensor/src/tests/ops/repeat.rs

@@ -45,4 +45,66 @@ mod tests {
  let data_expected = Data::from([[0, 1, 2], [0, 1, 2], [0, 1, 2], [0, 1, 2]]);
  assert_eq!(data_expected, data_actual);
  }
+
+ #[test]
+ fn should_support_float_repeat_on_dims_larger_than_1() {
+ let data = Data::from([
+ [[1.0, 2.0], [3.0, 4.0]],
+ [[5.0, 6.0], [7.0, 8.0]],
+ [[9.0, 10.0], [11.0, 12.0]],
+ [[13.0, 14.0], [15.0, 16.0]],
+ ]);
+ let tensor = Tensor::<TestBackend, 3>::from_data(data, &Default::default());
+
+ let data_actual = tensor.repeat(2, 2).into_data();
+
+ let data_expected = Data::from([
+ [[1.0, 2.0, 1.0, 2.0], [3.0, 4.0, 3.0, 4.0]],
+ [[5.0, 6.0, 5.0, 6.0], [7.0, 8.0, 7.0, 8.0]],
+ [[9.0, 10.0, 9.0, 10.0], [11.0, 12.0, 11.0, 12.0]],
+ [[13.0, 14.0, 13.0, 14.0], [15.0, 16.0, 15.0, 16.0]],
+ ]);
+
+ assert_eq!(data_expected, data_actual);
+ }
+
+ #[test]
+ fn should_support_int_repeat_on_dims_larger_than_1() {
+ let data = Data::from([
+ [[1, 2], [3, 4]],
+ [[5, 6], [7, 8]],
+ [[9, 10], [11, 12]],
+ [[13, 14], [15, 16]],
+ ]);
+ let tensor = Tensor::<TestBackend, 3, Int>::from_data(data, &Default::default());
+
+ let data_actual = tensor.repeat(2, 3).into_data();
+
+ let data_expected = Data::from([
+ [[1, 2, 1, 2, 1, 2], [3, 4, 3, 4, 3, 4]],
+ [[5, 6, 5, 6, 5, 6], [7, 8, 7, 8, 7, 8]],
+ [[9, 10, 9, 10, 9, 10], [11, 12, 11, 12, 11, 12]],
+ [[13, 14, 13, 14, 13, 14], [15, 16, 15, 16, 15, 16]],
+ ]);
+
+ assert_eq!(data_expected, data_actual);
+ }
+
+ #[test]
+ fn should_support_bool_repeat_on_dims_larger_than_1() {
+ let data = Data::from([
+ [[false, true], [true, false]],
+ [[true, true], [false, false]],
+ ]);
+ let tensor = Tensor::<TestBackend, 3, Bool>::from_data(data, &Default::default());
+
+ let data_actual = tensor.repeat(1, 2).into_data();
+
+ let data_expected = Data::from([
+ [[false, true], [true, false], [false, true], [true, false]],
+ [[true, true], [false, false], [true, true], [false, false]],
+ ]);
+
+ assert_eq!(data_expected, data_actual);
+ }
 }