move to moduleops

Author: crutcher

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

@@ -1,5 +1,4 @@
 use crate as burn;
-use burn_tensor::ops::linear::linear;
 
 use crate::config::Config;
 use crate::module::Param;

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

@@ -1,7 +1,6 @@
 use core::num::NonZeroUsize;
 
 use super::{conv, pool, unfold::unfold4d_using_conv2d};
-use crate::ops::linear::linear;
 use crate::{
     Shape, Tensor, TensorMetadata,
     backend::Backend,
@@ -796,7 +795,17 @@ pub trait ModuleOps<B: Backend> {
         weight: Tensor<B, 2>,
         bias: Option<Tensor<B, 1>>,
     ) -> Tensor<B, D> {
-        linear(input, weight, bias)
+        if D == 1 {
+            // Insert and remove an extra batch dimension for the batch matmul to work.
+            return Self::linear::<2>(input.unsqueeze(), weight, bias).flatten(0, 1);
+        }
+
+        let weight = weight.unsqueeze();
+        let output = input.matmul(weight);
+        match bias {
+            Some(bias) => output + bias.unsqueeze(),
+            None => output,
+        }
     }
 }
 

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

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