Extract Linear.forward to nn::functional::linear (#3147)

* Extract Linear.forward to nn::functional::linear

* remove linear_pytorch per review

* compat docs

* Refactor test assertions to use TensorData directly.

Replaced `Tensor::<TestBackend, _>::from_data` with `TensorData::from` for improved readability and consistency. Ensured data type conversion aligns with the input tensor's dtype. Added required imports to support changes.

* fmt

* move nn.functional to burn-tensor

* doc

* fix tests

* Lift linear to ModuleOps

* move to moduleops

* Refactor linear module ops with primitives

* Fix vector norm cuda tests / f16 tolerance

* Use alloc vec

* Update lock

---------

Co-authored-by: Guillaume Lagrange <[email protected]>

Author: crutcher

Cargo.lock

@@ -1,3 +1,5 @@
+use burn_tensor::module::linear;
+
 use crate as burn;
 
 use crate::config::Config;
@@ -24,7 +26,7 @@ pub struct LinearConfig {
     pub initializer: Initializer,
 }
 
-/// Applies a linear transformation to the input tensor:
+/// Applies a linear transformation to the input tensor.
 ///
 /// Should be created with [LinearConfig]
 ///
@@ -65,24 +67,24 @@ impl LinearConfig {
 impl<B: Backend> Linear<B> {
     /// Applies the forward pass on the input tensor.
     ///
+    /// # Arguments
+    ///
+    /// - `input` - The input tensor of shape `[..., d_input]`.
+    ///
     /// # Shapes
     ///
     /// - input: `[..., d_input]`
     /// - output: `[..., d_output]`
+    ///
+    /// # Returns
+    ///
+    /// The transformed tensor of shape `[..., d_output]`.
     pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
-        if D == 1 {
-            // Insert and remove an extra batch dimension for the batch matmul to work.
-            return Self::forward::<2>(self, input.unsqueeze()).flatten(0, 1);
-        }
-
-        let weight = self.weight.val().unsqueeze();
-        let bias = self.bias.as_ref().map(|b| b.val().unsqueeze());
-        let output = input.matmul(weight);
-
-        match bias {
-            Some(bias) => output + bias,
-            None => output,
-        }
+        linear(
+            input,
+            self.weight.val(),
+            self.bias.as_ref().map(|b| b.val()),
+        )
     }
 }
 

crates/burn-core/src/nn/linear.rs

@@ -357,3 +357,40 @@ where
         options,
     )))
 }
+
+/// Applies a [linear transformation](crate::ops::ModuleOps::linear) to the input tensor using the given weight and bias.
+///
+/// ```math
+/// y = x @ weight + [bias]
+/// ```
+///
+/// # Arguments:
+///
+/// - `input` is the input tensor, ``[..., d_input]``.
+/// - `weight` is the weight tensor, ``[d_input, d_output]``.
+/// - `bias` is the bias tensor (optional), ``[d_output]``.
+///
+/// # Returns:
+///
+/// The transformed tensor, ``[..., d_output]``.
+///
+/// # Compatibility
+///
+/// This function differs from PyTorch's ``torch.nn.functional.linear`` in that it does not
+/// transpose the weight matrix. In PyTorch, the weight matrix is transposed before
+/// multiplication:
+///
+/// ```math
+/// y = x @ weight^T + [bias]
+/// ```
+pub fn linear<B: Backend, const D: usize>(
+    input: Tensor<B, D>,
+    weight: Tensor<B, 2>,
+    bias: Option<Tensor<B, 1>>,
+) -> Tensor<B, D> {
+    Tensor::new(TensorPrimitive::Float(B::linear(
+        input.primitive.tensor(),
+        weight.primitive.tensor(),
+        bias.map(|b| b.primitive.tensor()),
+    )))
+}

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

@@ -1,3 +1,4 @@
+use alloc::vec;
 use core::num::NonZeroUsize;
 
 use super::{conv, pool, unfold::unfold4d_using_conv2d};
@@ -764,6 +765,66 @@ pub trait ModuleOps<B: Backend> {
         output_size: [usize; 2],
         options: InterpolateOptions,
     ) -> FloatTensor<B>;
+
+    /// Applies a linear transformation to the input tensor using the given weight and bias.
+    ///
+    /// ```math
+    /// y = x @ weight + [bias]
+    /// ```
+    ///
+    /// # Arguments:
+    ///
+    /// - `input` is the input tensor, ``[..., d_input]``.
+    /// - `weight` is the weight tensor, ``[d_input, d_output]``.
+    /// - `b` is the bias tensor (optional), ``[d_output]``.
+    ///
+    /// # Returns:
+    ///
+    /// The transformed tensor, ``[..., d_output]``.
+    ///
+    /// # Compatibility
+    ///
+    /// This function differs from PyTorch's ``torch.nn.functional.linear`` in that it does not
+    /// transpose the weight matrix. In PyTorch, the weight matrix is transposed before
+    /// multiplication:
+    ///
+    /// ```math
+    /// y = x @ weight^T + [bias]
+    /// ```
+    fn linear(
+        input: FloatTensor<B>,
+        weight: FloatTensor<B>,
+        bias: Option<FloatTensor<B>>,
+    ) -> FloatTensor<B> {
+        let ndims_in = input.shape().num_dims();
+        let [d_input, d_output] = weight.shape().dims();
+        if ndims_in == 1 {
+            // Insert and remove an extra batch dimension for the batch matmul to work.
+            let input = B::float_reshape(input, Shape::from([1, d_input]));
+            let output = Self::linear(input, weight, bias);
+            return B::float_reshape(output, Shape::from([d_output]));
+        }
+
+        let weight = unsqueeze::<B>(weight, ndims_in);
+        let output = B::float_matmul(input, weight);
+        match bias {
+            Some(bias) => B::float_add(output, unsqueeze::<B>(bias, ndims_in)),
+            None => output,
+        }
+    }
+}
+
+// Unsqueeze op on primitive.
+// TODO: would be nice to have this on primitives too for convenience.
+fn unsqueeze<B: Backend>(tensor: FloatTensor<B>, ndims_out: usize) -> FloatTensor<B> {
+    let shape = tensor.shape();
+    let ndims_in = shape.num_dims();
+
+    let mut dims = vec![1; ndims_out];
+    let num_ones = ndims_out - ndims_in;
+    dims[num_ones..(ndims_in + num_ones)].copy_from_slice(&shape.dims[..ndims_in]);
+
+    B::float_reshape(tensor, Shape::from(dims))
 }
 
 #[cfg(test)]