fix tests

Author: crutcher

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

@@ -1,5 +1,4 @@
-use crate::tensor::Tensor;
-use crate::tensor::backend::Backend;
+use crate::{Tensor, backend::Backend};
 
 /// Applies a linear transformation to the input tensor using the given weight and bias.
 ///
@@ -43,67 +42,3 @@ pub fn linear<B: Backend, const D: usize>(
         None => output,
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-    use crate::TestBackend;
-    use burn_tensor::TensorData;
-
-    #[test]
-    fn test_linear_1d() {
-        let weight =
-            Tensor::<TestBackend, 2>::from_data([[1.0, 2.0], [3.0, 4.0]], &Default::default());
-
-        let x = Tensor::<TestBackend, 1>::from_data([1.0, 2.0], &Default::default());
-
-        linear(x.clone(), weight.clone(), None)
-            .into_data()
-            .assert_eq(
-                &TensorData::from([7.0, 10.0]).convert_dtype(x.dtype()),
-                true,
-            );
-    }
-
-    #[test]
-    fn test_linear_forward_no_bias() {
-        let weight =
-            Tensor::<TestBackend, 2>::from_data([[1.0, 2.0], [3.0, 4.0]], &Default::default());
-
-        let x = Tensor::<TestBackend, 3>::from_data(
-            [[[1.0, 2.0], [3.0, 4.0]], [[-1.0, -2.0], [-3.0, -4.0]]],
-            &Default::default(),
-        );
-
-        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]]])
-                    .convert_dtype(x.dtype()),
-                true,
-            );
-    }
-
-    #[test]
-    fn test_linear_forward_with_bias() {
-        let weight =
-            Tensor::<TestBackend, 2>::from_data([[1.0, 2.0], [3.0, 4.0]], &Default::default());
-        let bias = Some(Tensor::<TestBackend, 1>::from_data(
-            [1.0, -1.0],
-            &Default::default(),
-        ));
-
-        let x = Tensor::<TestBackend, 3>::from_data(
-            [[[1.0, 2.0], [3.0, 4.0]], [[-1.0, -2.0], [-3.0, -4.0]]],
-            &Default::default(),
-        );
-
-        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]]])
-                    .convert_dtype(x.dtype()),
-                true,
-            );
-    }
-}

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

@@ -2,6 +2,7 @@ mod activation;
 mod clone_invariance;
 mod grid;
 mod module;
+mod nn;
 mod ops;
 mod primitive;
 mod quantization;
@@ -155,6 +156,9 @@ 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!();
 

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

@@ -0,0 +1 @@
+pub(crate) mod vector_norm;

crates/burn-tensor/src/tests/nn/functional/vector_norm.rs

@@ -0,0 +1,51 @@
+#[burn_tensor_testgen::testgen(nn_fn_vector_norm)]
+mod tests {
+    use super::*;
+    use burn_tensor::nn::functional::linear;
+    use burn_tensor::{Tensor, TensorData};
+
+    #[test]
+    fn test_linear_1d() {
+        let weight = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]);
+
+        let x = TestTensor::<1>::from([1.0, 2.0]);
+
+        linear(x.clone(), weight.clone(), None)
+            .into_data()
+            .assert_eq(
+                &TensorData::from([7.0, 10.0]).convert_dtype(x.dtype()),
+                true,
+            );
+    }
+
+    #[test]
+    fn test_linear_forward_no_bias() {
+        let weight = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]);
+
+        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)
+            .into_data()
+            .assert_eq(
+                &TensorData::from([[[7.0, 10.0], [15.0, 22.0]], [[-7.0, -10.0], [-15.0, -22.0]]])
+                    .convert_dtype(x.dtype()),
+                true,
+            );
+    }
+
+    #[test]
+    fn test_linear_forward_with_bias() {
+        let weight = TestTensor::<2>::from([[1.0, 2.0], [3.0, 4.0]]);
+        let bias = Some(TestTensor::<1>::from([1.0, -1.0]));
+
+        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())
+            .into_data()
+            .assert_eq(
+                &TensorData::from([[[8.0, 9.0], [16.0, 21.0]], [[-6.0, -11.0], [-14.0, -23.0]]])
+                    .convert_dtype(x.dtype()),
+                true,
+            );
+    }
+}

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

@@ -0,0 +1 @@
+pub(crate) mod functional;