Lift linear to ModuleOps

Author: crutcher

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

@@ -1,9 +1,9 @@
 use crate as burn;
+use burn_tensor::ops::linear::linear;
 
 use crate::config::Config;
 use crate::module::Param;
 use crate::module::{Content, DisplaySettings, Module, ModuleDisplay};
-use crate::nn::functional::linear;
 use crate::tensor::{Tensor, backend::Backend};
 
 use super::Initializer;
@@ -30,8 +30,6 @@ pub struct LinearConfig {
 /// Should be created with [LinearConfig]
 ///
 /// `O = IW + b`
-///
-/// See: [linear][nn::functional::linear]
 #[derive(Module, Debug)]
 #[module(custom_display)]
 pub struct Linear<B: Backend> {
@@ -83,7 +81,8 @@ impl<B: Backend> Linear<B> {
     pub fn forward<const D: usize>(&self, input: Tensor<B, D>) -> Tensor<B, D> {
         let weight = self.weight.val();
         let bias = self.bias.as_ref().map(|b| b.val());
-        linear(input, weight, bias)
+
+        B::linear(input, weight, bias)
     }
 }
 

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

@@ -1,5 +1,3 @@
-pub use burn_tensor::nn::*;
-
 /// Attention module
 pub mod attention;
 

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

@@ -32,9 +32,6 @@ pub mod loss;
 /// The burn module.
 pub mod module;
 
-/// The nn module.
-pub mod nn;
-
 /// Operations on tensors module.
 pub mod ops;
 

crates/burn-tensor/src/tensor/nn/functional/mod.rs

@@ -1,8 +1,9 @@
 use core::num::NonZeroUsize;
 
 use super::{conv, pool, unfold::unfold4d_using_conv2d};
+use crate::ops::linear::linear;
 use crate::{
-    Shape, TensorMetadata,
+    Shape, Tensor, TensorMetadata,
     backend::Backend,
     ops::{FloatTensor, IntTensor},
 };
@@ -764,6 +765,39 @@ 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<const D: usize>(
+        input: Tensor<B, D>,
+        weight: Tensor<B, 2>,
+        bias: Option<Tensor<B, 1>>,
+    ) -> Tensor<B, D> {
+        linear(input, weight, bias)
+    }
 }
 
 #[cfg(test)]

crates/burn-tensor/src/tensor/nn/mod.rs

@@ -13,4 +13,7 @@ pub mod pool;
 
 mod base;
 
+/// Module with linear operations.
+pub mod linear;
+
 pub use base::*;

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

@@ -2,7 +2,6 @@ mod activation;
 mod clone_invariance;
 mod grid;
 mod module;
-mod nn;
 mod ops;
 mod primitive;
 mod quantization;
@@ -156,9 +155,6 @@ macro_rules! testgen_with_float_param {
         burn_tensor::testgen_silu!();
         burn_tensor::testgen_tanh_activation!();
 
-        // test nn.functional
-        burn_tensor::testgen_nn_fn_vector_norm!();
-
         // test grid
         burn_tensor::testgen_meshgrid!();
 
@@ -181,6 +177,7 @@ macro_rules! testgen_with_float_param {
         burn_tensor::testgen_module_nearest_interpolate!();
         burn_tensor::testgen_module_bilinear_interpolate!();
         burn_tensor::testgen_module_bicubic_interpolate!();
+        burn_tensor::testgen_module_linear!();
 
         // test ops
         burn_tensor::testgen_gather_scatter!();

crates/burn-tensor/src/tensor/nn/functional/linear.rs renamed to crates/burn-tensor/src/tensor/ops/modules/linear.rs

@@ -1,7 +1,7 @@
-#[burn_tensor_testgen::testgen(nn_fn_vector_norm)]
+#[burn_tensor_testgen::testgen(module_linear)]
 mod tests {
     use super::*;
-    use burn_tensor::nn::functional::linear;
+    use burn_tensor::ops::ModuleOps;
     use burn_tensor::{Tensor, TensorData};
 
     #[test]
@@ -10,7 +10,7 @@ mod tests {
 
         let x = TestTensor::<1>::from([1.0, 2.0]);
 
-        linear(x.clone(), weight.clone(), None)
+        TestBackend::linear(x.clone(), weight.clone(), None)
             .into_data()
             .assert_eq(
                 &TensorData::from([7.0, 10.0]).convert_dtype(x.dtype()),
@@ -24,7 +24,7 @@ mod tests {
 
         let x = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[-1.0, -2.0], [-3.0, -4.0]]]);
 
-        linear(x.clone(), weight.clone(), None)
+        TestBackend::linear(x.clone(), weight.clone(), None)
             .into_data()
             .assert_eq(
                 &TensorData::from([[[7.0, 10.0], [15.0, 22.0]], [[-7.0, -10.0], [-15.0, -22.0]]])
@@ -40,7 +40,7 @@ mod tests {
 
         let x = TestTensor::<3>::from([[[1.0, 2.0], [3.0, 4.0]], [[-1.0, -2.0], [-3.0, -4.0]]]);
 
-        linear(x.clone(), weight.clone(), bias.clone())
+        TestBackend::linear(x.clone(), weight.clone(), bias.clone())
             .into_data()
             .assert_eq(
                 &TensorData::from([[[8.0, 9.0], [16.0, 21.0]], [[-6.0, -11.0], [-14.0, -23.0]]])

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

@@ -30,6 +30,7 @@ mod full;
 mod gather_scatter;
 mod init;
 mod iter_dim;
+pub(crate) mod linear;
 mod log;
 mod log1p;
 mod map_comparison;