Tensor expand operator

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

@@ -144,6 +144,7 @@ Those operations are available for all tensor kinds: `Int`, `Float`, and `Bool`.
 | `tensor.all_dim(dim)`                 | `tensor.all(dim)`                    |
 | `tensor.any()`                        | `tensor.any()`                       |
 | `tensor.any_dim(dim)`                 | `tensor.any(dim)`                    |
+| `tensor.expand(shape)`                | `tensor.expand(shape)`               |
 | `tensor.chunk(num_chunks, dim)`       | `tensor.chunk(num_chunks, dim)`      |
 | `tensor.device()`                     | `tensor.device`                      |
 | `tensor.dims()`                       | `tensor.size()`                      |

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

@@ -128,4 +128,11 @@ impl<B: Backend, C: CheckpointStrategy> BoolTensorOps<Self> for Autodiff<B, C> {
     fn bool_nonzero<const D: usize>(tensor: BoolTensor<B, D>) -> Vec<IntTensor<B, 1>> {
         B::bool_nonzero(tensor)
     }
+
+    fn bool_expand<const D: usize, const D2: usize>(
+        tensor: BoolTensor<B, D>,
+        shape: Shape<D2>,
+    ) -> BoolTensor<B, D2> {
+        B::bool_expand(tensor, shape)
+    }
 }

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

@@ -369,6 +369,13 @@ impl<B: Backend, C: CheckpointStrategy> IntTensorOps<Self> for Autodiff<B, C> {
         B::int_prod_dim(tensor, dim)
     }
 
+    fn int_expand<const D: usize, const D2: usize>(
+        tensor: IntTensor<B, D>,
+        shape: Shape<D2>,
+    ) -> IntTensor<B, D2> {
+        B::int_expand(tensor, shape)
+    }
+
     fn int_sort<const D: usize>(
         tensor: IntTensor<Self, D>,
         dim: usize,

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

@@ -2437,6 +2437,81 @@ impl<B: Backend, C: CheckpointStrategy> FloatTensorOps<Self> for Autodiff<B, C>
             .stateless(B::float_sign(tensor.primitive))
     }
 
+    fn float_expand<const D1: usize, const D2: usize>(
+        tensor: FloatTensor<Self, D1>,
+        shape: Shape<D2>,
+    ) -> FloatTensor<Self, D2> {
+        #[derive(Debug)]
+        struct ExpandDim<const D1: usize, const D2: usize>;
+
+        #[derive(new, Debug)]
+        struct RetroExpand<B: Backend, const D1: usize, const D2: usize> {
+            input_id: NodeID,
+            shape: Shape<D2>,
+            _backend: PhantomData<B>,
+        }
+
+        impl<B: Backend, const D1: usize, const D2: usize> RetroForward for RetroExpand<B, D1, D2> {
+            fn forward(&self, states: &mut BackwardStates, out_node: NodeID) {
+                let input = states.get_state::<B::FloatTensorPrimitive<D1>>(&self.input_id);
+                let out = B::float_expand(input, self.shape.clone());
+                states.save(out_node, out)
+            }
+        }
+
+        impl<B: Backend, const D1: usize, const D2: usize> Backward<B, D2, 1> for ExpandDim<D1, D2> {
+            type State = Shape<D1>;
+
+            fn backward(
+                self,
+                ops: Ops<Self::State, 1>,
+                grads: &mut Gradients,
+                _checkpointer: &mut Checkpointer,
+            ) {
+                let shape_original = ops.state;
+
+                let mut shape_expanded = [1; D2];
+
+                debug_assert!(D2 >= D1);
+
+                for i in 0..D1 {
+                    shape_expanded[i + (D2 - D1)] = shape_original.dims[i];
+                }
+
+                unary::<B, D2, D1, _>(ops.parents, ops.node, grads, |grad| {
+                    let shape_grad = B::float_shape(&grad);
+                    let mut grad = grad;
+
+                    #[allow(clippy::needless_range_loop)]
+                    for i in 0..D2 {
+                        if shape_expanded[i] == 1 && shape_grad.dims[i] != 1 {
+                            grad = B::float_sum_dim(grad, i);
+                        }
+                    }
+
+                    B::float_reshape(grad, shape_original)
+                });
+            }
+        }
+
+        match ExpandDim
+            .prepare::<C>([tensor.node.clone()], [tensor.graph.clone()])
+            .memory_bound()
+            .retro_forward(RetroExpand::<B, D1, D2>::new(
+                tensor.node.id.clone(),
+                shape.clone(),
+            ))
+            .parents([&tensor])
+            .stateful()
+        {
+            OpsKind::Tracked(prep) => prep.finish(
+                B::float_shape(&tensor.primitive),
+                B::float_expand(tensor.primitive, shape),
+            ),
+            OpsKind::UnTracked(prep) => prep.finish(B::float_expand(tensor.primitive, shape)),
+        }
+    }
+
     fn float_sort<const D: usize>(
         tensor: FloatTensor<Self, D>,
         dim: usize,

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

@@ -0,0 +1,38 @@
+#[burn_tensor_testgen::testgen(ad_expand)]
+mod tests {
+    use super::*;
+    use burn_tensor::{Data, Tensor};
+
+    #[test]
+    fn should_diff_expand() {
+        // Python code to generate the test case values
+        // import torch
+        // x1 = torch.tensor([4.0, 7.0, 2.0, 3.0], requires_grad=True)
+        // x2 = torch.tensor([2.0, 4.5, 7.0, 3.0], requires_grad=True)
+        // y = x1.expand(4, 4)
+        // z = (x2 * y).sum()
+        // z.backward()
+        // print("x1", x1.grad)
+        // print("x2", x2.grad)
+
+        let device = Default::default();
+
+        let data_1: Data<f32, 1> = Data::from([4.0, 7.0, 2.0, 3.0]);
+        let tensor_1 = TestAutodiffTensor::from_data(data_1, &device).require_grad();
+
+        let data_2: Data<f32, 1> = Data::from([2.0, 4.5, 7.0, 3.0]);
+        let tensor_2 = TestAutodiffTensor::from_data(data_2, &device).require_grad();
+
+        let tensor_3 = tensor_1.clone().expand([4, 4]);
+
+        // Use unsqueeze to make tensor_2 have the same shape as tensor_3
+        let tensor_4 = tensor_2.clone().unsqueeze().mul(tensor_3).sum();
+        let grads = tensor_4.backward();
+
+        let grad_1 = tensor_1.grad(&grads).unwrap();
+        let grad_2 = tensor_2.grad(&grads).unwrap();
+
+        assert_eq!(grad_1.to_data(), Data::from([8., 18., 28., 12.]));
+        assert_eq!(grad_2.to_data(), Data::from([16., 28., 8., 12.]));
+    }
+}

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

@@ -21,6 +21,7 @@ mod cross_entropy;
 mod div;
 mod erf;
 mod exp;
+mod expand;
 mod flip;
 mod gather_scatter;
 mod gelu;
@@ -119,6 +120,7 @@ macro_rules! testgen_all {
         burn_autodiff::testgen_ad_flip!();
         burn_autodiff::testgen_ad_nonzero!();
         burn_autodiff::testgen_ad_sign!();
+        burn_autodiff::testgen_ad_expand!();
         burn_autodiff::testgen_ad_sort!();
     };
 }

crates/burn-candle/src/lib.rs

@@ -102,6 +102,7 @@ mod tests {
     burn_tensor::testgen_sub!();
     burn_tensor::testgen_tanh!();
     burn_tensor::testgen_transpose!();
+    burn_tensor::testgen_expand!();
 
     // test stats
     burn_tensor::testgen_var!();
@@ -157,4 +158,5 @@ mod tests {
     burn_autodiff::testgen_ad_sub!();
     burn_autodiff::testgen_ad_tanh!();
     burn_autodiff::testgen_ad_transpose!();
+    burn_autodiff::testgen_ad_expand!();
 }

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

@@ -142,3 +142,10 @@ pub fn chunk<E: CandleElement, const D: usize>(
         Err(e) => panic!("error chunk from Candle"),
     }
 }
+
+pub fn expand<E: CandleElement, const D1: usize, const D2: usize>(
+    tensor: CandleTensor<E, D1>,
+    shape: Shape<D2>,
+) -> CandleTensor<E, D2> {
+    CandleTensor::new(tensor.tensor.broadcast_as(&shape.dims).unwrap())
+}

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

@@ -8,6 +8,8 @@ use crate::{
     Candle, CandleTensor,
 };
 
+use super::base::{expand, permute};
+
 impl<F: FloatCandleElement, I: IntCandleElement> BoolTensorOps<Self> for Candle<F, I> {
     fn bool_empty<const D: usize>(shape: Shape<D>, device: &Device<Self>) -> BoolTensor<Self, D> {
         super::base::empty(shape, device)
@@ -140,4 +142,11 @@ impl<F: FloatCandleElement, I: IntCandleElement> BoolTensorOps<Self> for Candle<
     ) -> BoolTensor<Self, D> {
         super::base::flip(tensor, axes)
     }
+
+    fn bool_expand<const D1: usize, const D2: usize>(
+        tensor: BoolTensor<Self, D1>,
+        shape: Shape<D2>,
+    ) -> BoolTensor<Self, D2> {
+        expand(tensor, shape)
+    }
 }

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

@@ -8,6 +8,8 @@ use crate::{
     Candle, CandleTensor,
 };
 
+use super::base::{expand, permute};
+
 impl<F: FloatCandleElement, I: IntCandleElement> IntTensorOps<Self> for Candle<F, I> {
     fn int_empty<const D: usize>(shape: Shape<D>, device: &Device<Self>) -> IntTensor<Self, D> {
         super::base::empty(shape, device)
@@ -430,6 +432,13 @@ impl<F: FloatCandleElement, I: IntCandleElement> IntTensorOps<Self> for Candle<F
         super::base::flip(tensor, axes)
     }
 
+    fn int_expand<const D1: usize, const D2: usize>(
+        tensor: IntTensor<Self, D1>,
+        shape: Shape<D2>,
+    ) -> IntTensor<Self, D2> {
+        expand(tensor, shape)
+    }
+
     // TODO add sign operator once Candle supports it:
     // https://github.com/huggingface/candle/issues/1827
 }

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

@@ -11,6 +11,8 @@ use crate::{
     Candle, CandleTensor,
 };
 
+use super::base::{expand, permute};
+
 impl<F: FloatCandleElement, I: IntCandleElement> FloatTensorOps<Self> for Candle<F, I> {
     fn float_from_data<const D: usize>(
         data: Data<F, D>,
@@ -530,6 +532,13 @@ impl<F: FloatCandleElement, I: IntCandleElement> FloatTensorOps<Self> for Candle
         super::base::flip(tensor, axes)
     }
 
+    fn float_expand<const D1: usize, const D2: usize>(
+        tensor: FloatTensor<Self, D1>,
+        shape: Shape<D2>,
+    ) -> FloatTensor<Self, D2> {
+        expand(tensor, shape)
+    }
+
     // TODO add sign operator once Candle supports it:
     // https://github.com/huggingface/candle/issues/1827
 }

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

@@ -4,9 +4,10 @@ use crate::{
     ops::binary::binary_ops_shape,
     stream::{
         BaseOperationDescription, BinaryOperationDescription, BoolOperationDescription,
-        CatOperationDescription, FlipOperationDescription, Operation, OperationDescription,
-        PermuteOperationDescription, ReshapeDescription, SliceAssignOperationDescription,
-        SliceOperationDescription, StreamId, SwapDimsDescription, UnaryOperationDescription,
+        CatOperationDescription, ExpandOperationDescription, FlipOperationDescription, Operation,
+        OperationDescription, PermuteOperationDescription, ReshapeDescription,
+        SliceAssignOperationDescription, SliceOperationDescription, StreamId, SwapDimsDescription,
+        UnaryOperationDescription,
     },
     Fusion, FusionBackend,
 };
@@ -467,6 +468,44 @@ impl<B: FusionBackend> BoolTensorOps<Self> for Fusion<B> {
         out
     }
 
+    fn bool_expand<const D1: usize, const D2: usize>(
+        tensor: BoolTensor<Self, D1>,
+        shape: Shape<D2>,
+    ) -> BoolTensor<Self, D2> {
+        #[derive(new)]
+        struct ExpandOps<const D: usize, const D2: usize> {
+            desc: ExpandOperationDescription,
+        }
+
+        impl<const D: usize, const D2: usize, B: FusionBackend> Operation<B> for ExpandOps<D, D2> {
+            fn execute(self: Box<Self>, handles: &mut crate::HandleContainer<B>) {
+                let input = handles.get_bool_tensor::<D>(&self.desc.input);
+                let shape: [usize; D2] = self.desc.shape.try_into().unwrap();
+                let output = B::bool_expand(input, shape.into());
+
+                handles.register_bool_tensor(&self.desc.out.id, output);
+            }
+        }
+
+        let stream = tensor.stream;
+
+        let out = tensor.client.tensor_uninitialized(shape.dims.into());
+
+        let desc = ExpandOperationDescription {
+            input: tensor.into_description(),
+            shape: shape.dims.into(),
+            out: out.to_description_out(),
+        };
+
+        out.client.register(
+            vec![stream],
+            OperationDescription::BaseBool(BaseOperationDescription::Expand(desc.clone())),
+            ExpandOps::<D1, D2>::new(desc),
+        );
+
+        out
+    }
+
     fn bool_flip<const D: usize>(
         tensor: BoolTensor<Self, D>,
         axes: &[usize],

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

@@ -6,13 +6,14 @@ use crate::{
     scalar_float2int_ops, scalar_float_cmp_ops, scalar_float_ops,
     stream::{
         BaseOperationDescription, BinaryOperationDescription, CatOperationDescription,
-        ClampOperationDescription, FlipOperationDescription, FloatOperationDescription,
-        GatherOperationDescription, MaskFillOperationDescription, MaskWhereOperationDescription,
-        NumericOperationDescription, Operation, OperationDescription, PermuteOperationDescription,
-        RandomOperationDescription, ReduceDimWithIndicesDescription, ReshapeDescription,
-        ScalarOperationDescription, ScatterOperationDescription, SelectAssignOperationDescription,
-        SelectOperationDescription, SliceAssignOperationDescription, SliceOperationDescription,
-        StreamId, SwapDimsDescription, UnaryOperationDescription,
+        ClampOperationDescription, ExpandOperationDescription, FlipOperationDescription,
+        FloatOperationDescription, GatherOperationDescription, MaskFillOperationDescription,
+        MaskWhereOperationDescription, NumericOperationDescription, Operation,
+        OperationDescription, PermuteOperationDescription, RandomOperationDescription,
+        ReduceDimWithIndicesDescription, ReshapeDescription, ScalarOperationDescription,
+        ScatterOperationDescription, SelectAssignOperationDescription, SelectOperationDescription,
+        SliceAssignOperationDescription, SliceOperationDescription, StreamId, SwapDimsDescription,
+        UnaryOperationDescription,
     },
     unary_float_ops, Fusion, FusionBackend, TensorDescription,
 };
@@ -1847,6 +1848,44 @@ impl<B: FusionBackend> FloatTensorOps<Self> for Fusion<B> {
         out
     }
 
+    fn float_expand<const D1: usize, const D2: usize>(
+        tensor: FloatTensor<Self, D1>,
+        shape: Shape<D2>,
+    ) -> FloatTensor<Self, D2> {
+        #[derive(new)]
+        struct ExpandOps<const D: usize, const D2: usize> {
+            desc: ExpandOperationDescription,
+        }
+
+        impl<const D: usize, const D2: usize, B: FusionBackend> Operation<B> for ExpandOps<D, D2> {
+            fn execute(self: Box<Self>, handles: &mut crate::HandleContainer<B>) {
+                let input = handles.get_float_tensor::<D>(&self.desc.input);
+                let shape: [usize; D2] = self.desc.shape.try_into().unwrap();
+                let output = B::float_expand(input, shape.into());
+
+                handles.register_float_tensor(&self.desc.out.id, output);
+            }
+        }
+
+        let stream = tensor.stream;
+
+        let out = tensor.client.tensor_uninitialized(shape.dims.into());
+
+        let desc = ExpandOperationDescription {
+            input: tensor.into_description(),
+            shape: shape.dims.into(),
+            out: out.to_description_out(),
+        };
+
+        out.client.register(
+            vec![stream],
+            OperationDescription::BaseFloat(BaseOperationDescription::Expand(desc.clone())),
+            ExpandOps::<D1, D2>::new(desc),
+        );
+
+        out
+    }
+
     fn float_flip<const D: usize>(
         tensor: FloatTensor<Self, D>,
         axes: &[usize],