Updates

Author: catch-twenty-two

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

@@ -5,83 +5,78 @@ use burn::{prelude::Backend, tensor::Tensor};
 /// ```ignore
 /// broadcast!(
 ///     a: Tensor<Backend, RANK_A>,
-///     b: Tensor<Backend, RANK_B>
+///     b: Tensor<RANK_B>
 /// )
 /// ```
 ///
 /// # Parameters
 /// - `a`: Identifier for the first tensor variable (e.g., `a`).
-/// - `backend`: The backend type used in the first tensor (e.g., `MyBackend`).
-/// - `dims1`: The static rank of the first tensor (e.g., `2`, `3`, etc.).
+/// - `Backend`: The backend to use
+/// - `RANK_A`: The static rank of the first tensor (e.g., `2`, `3`, etc.).
 ///
 /// - `b`: Identifier for the second tensor variable (e.g., `b`).
-/// - `backend`: The backend type used in the second tensor (must match `backend` for correctness).
-/// - `dims2`: The static rank of the second tensor.
-///
-/// # Expansion
-/// Expands to:
-/// ```rust
-/// {
-///     const N: usize = max(dims1, dims2);
-///     broadcast::<B, N, dims1, dims2>(a, b)
-/// }
-/// ```
+/// - `RANK_B`: The static rank of the second tensor.
 ///
 /// # Example
 /// ```rust
-/// let a: Tensor<MyBackend, 2> = ...;
-/// let b: Tensor<MyBackend, 4> = ...;
+///     let device = &NdArrayDevice::default();
+///     type B = NdArray<f32>;
 ///
-/// let result = broadcast!(
-///     a: Tensor<MyBackend, 2>,
-///     b: Tensor<MyBackend, 4>
-/// );
-/// // Expands to: broadcast::<MyBackend, 4, 2, 4>(a, b)
-/// ```
+///     let a = Tensor::<B, 3>::from_data(
+///         [
+///             [[2, 8, 7, 2], [9, 14, 13, 12], [9, 14, 13, 12]],
+///             [[2, 8, 7, 2], [9, 14, 13, 12], [9, 14, 13, 12]],
+///         ],
+///         device,
+///     );
 ///
+///     let b = Tensor::<B, 2>::from_data([[4, 11, 10, 5]], device);
+///
+///     let (a, b) = broadcast!(a:Tensor<B, 3>, b:Tensor<2>);
+///
+///     let a_add_b = a.add(b);
+///
+/// // Output:
+/// // Tensor {
+/// //   data:
+/// // [[[ 6.0, 19.0, 17.0,  7.0],
+/// //   [13.0, 25.0, 23.0, 17.0],
+/// //   [13.0, 25.0, 23.0, 17.0]],
+/// //  [[ 6.0, 19.0, 17.0,  7.0],
+/// //   [13.0, 25.0, 23.0, 17.0],
+/// //   [13.0, 25.0, 23.0, 17.0]]],
+/// //   shape:  [2, 3, 4],
+/// //   device:  Cpu,
+/// //   backend:  "ndarray",
+/// //   kind:  "Float",
+/// //   dtype:  "f32",
+/// // }
+/// ```
 #[macro_export]
 macro_rules! broadcast {
     (
-        $a:ident : Tensor<$backend1:ty, $dims1:tt>,
-        $b:ident : Tensor<$backend2:ty, $dims2:tt>
+        $a:ident : Tensor<$backend:ty, $dims1:tt>,
+        $b:ident : Tensor<$dims2:tt>
     ) => {{
-        use $crate::ops::broadcast_op;
+        use $crate::broadcast::broadcast_op;
         const fn max(a: usize, b: usize) -> usize {
             if a > b { a } else { b }
         }
-    
-        const N: usize = max($dims1, $dims2);
-
-        broadcast_op::<$backend1, N, $dims1, $dims2>($a, $b)
-    }};
-}
 
-#[macro_export]
-macro_rules! add_broadcast {
-    (
-        $a:ident : Tensor<$backend1:ty, $dims1:tt>,
-        $b:ident : Tensor<$backend2:ty, $dims2:tt>
-    ) => {{
-        use $crate::ops::broadcast_op;
-        const fn max(a: usize, b: usize) -> usize {
-            if a > b { a } else { b }
-        }
-    
         const N: usize = max($dims1, $dims2);
 
-        let (a,b) = broadcast_op::<$backend1, N, $dims1, $dims2>($a, $b);
-        a.add(b)
+        broadcast_op::<$backend, N, $dims1, $dims2>(&$a, &$b)
     }};
 }
 
 pub fn broadcast_op<B: Backend, const N: usize, const DA: usize, const DB: usize>(
-    a: Tensor<B, DA>,
-    b: Tensor<B, DB>,
+    a: &Tensor<B, DA>,
+    b: &Tensor<B, DB>,
 ) -> (Tensor<B, N>, Tensor<B, N>) {
     // pad left with 1s
 
-    let a = a.unsqueeze::<N>();
-    let b = b.unsqueeze::<N>();
+    let a = a.clone().unsqueeze::<N>();
+    let b = b.clone().unsqueeze::<N>();
 
     let b_shape = b.shape().dims::<N>();
 
@@ -129,8 +124,8 @@ pub fn broadcast_op<B: Backend, const N: usize, const DA: usize, const DB: usize
 
 #[cfg(test)]
 mod tests {
-    use burn::backend::ndarray::{NdArray, NdArrayDevice};
     use super::*;
+    use burn::backend::ndarray::{NdArray, NdArrayDevice};
 
     #[test]
     fn test_broadcast_multi_dims() {
@@ -140,7 +135,7 @@ mod tests {
         let a = Tensor::<B, 6>::empty([7, 6, 2, 3, 1, 9], device);
         let b = Tensor::<B, 4>::empty([2, 1, 7, 1], device);
 
-        let (a, b) = broadcast!(a: Tensor<B, 6>, b: Tensor<B, 4>);
+        let (a, b) = broadcast!(a: Tensor<B, 6>, b: Tensor<4>);
 
         assert_eq!(a.shape(), b.shape());
     }
@@ -160,14 +155,21 @@ mod tests {
 
         let b = Tensor::<B, 2>::from_data([[4, 11, 10, 5]], device);
 
-        let (a, b) = broadcast!(a:Tensor<B, 3>, b:Tensor<B, 2>);
-
+        let (a, b) = broadcast!(a:Tensor<B, 3>, b:Tensor<2>);
         let a_add_b = a.add(b);
 
         Tensor::<B, 3>::from_data(
             [
-                [[6, 19, 17, 7], [13, 25, 23, 17], [13, 25, 23, 17]],
-                [[6, 19, 17, 7], [13, 25, 23, 17], [13, 25, 23, 17]],
+                [
+                    [6.0, 19.0, 17.0, 7.0],
+                    [13.0, 25.0, 23.0, 17.0],
+                    [13.0, 25.0, 23.0, 17.0],
+                ],
+                [
+                    [6.0, 19.0, 17.0, 7.0],
+                    [13.0, 25.0, 23.0, 17.0],
+                    [13.0, 25.0, 23.0, 17.0],
+                ],
             ],
             device,
         )
@@ -181,13 +183,10 @@ mod tests {
         type B = NdArray<f32>;
 
         let a = Tensor::<B, 1>::from_data([3.0, 2.0, 6.0, 3.0], device);
-
         let b = Tensor::<B, 1>::from_data([1.0, 0.5, 4.0, 7.0], device);
-
         let a = a.reshape([-1, 1]);
 
-        let (a, b) = broadcast!(a:Tensor<B, 2>, b:Tensor<B, 1>);
-
+        let (a, b) = broadcast!(a:Tensor<B, 2>, b:Tensor<1>);
         let max_a_b = a.max_pair(b);
 
         Tensor::<B, 2>::from_data(
@@ -209,12 +208,10 @@ mod tests {
         type B = NdArray<f32>;
 
         let a = Tensor::<B, 1>::from_data([1.1, 2.2, 3.3], device);
-
         let b = Tensor::<B, 1>::from_data([4.0, 5.0, 6.0, 7.0], device);
-
         let a = a.reshape([-1, 1]);
 
-        let (a, b) = broadcast!(a:Tensor<B, 2>, b:Tensor<B, 1>);
+        let (a, b) = broadcast!(a:Tensor<B, 2>, b:Tensor<1>);
         let add_a_b = a.add(b);
 
         Tensor::<B, 2>::from_data(
@@ -232,7 +229,7 @@ mod tests {
         let b = Tensor::<B, 1>::from_data([4.0, 5.0, 6.0, 7.0], device);
 
         let b = b.reshape([-1, 1]);
-        let (a, b) = broadcast!(a:Tensor<B, 1>, b:Tensor<B, 2>);
+        let (a, b) = broadcast!(a:Tensor<B, 1>, b:Tensor<2>);
         let add_a_b = a.add(b);
 
         Tensor::<B, 2>::from_data(
@@ -254,13 +251,10 @@ mod tests {
         type B = NdArray<f32>;
 
         let a = Tensor::<B, 1>::from_data([3.0, 2.0, 6.0, 3.0], device);
-
         let b = Tensor::<B, 1>::from_data([1.0, 0.5, 4.0, 7.0, 8.0], device);
 
         let b = b.reshape([-1, 1]);
-
-        let (a, b) = broadcast!(a:Tensor<B, 1>, b:Tensor<B, 2>);
-
+        let (a, b) = broadcast!(a:Tensor<B, 1>, b:Tensor<2>);
         let max_a_b = a.max_pair(b);
 
         Tensor::<B, 2>::from_data(
@@ -292,7 +286,7 @@ mod tests {
         );
 
         let b = Tensor::<B, 1>::from_data([1.0, 0.5, 4.0, 7.0], device);
-        let (a, b) = broadcast!(a:Tensor<B, 2>, b:Tensor<B, 1>);
+        let (a, b) = broadcast!(a:Tensor<B, 2>, b:Tensor<1>);
 
         let add_a_b = a.add(b);