Refactor qtensor ops tests w/ QTensor helper struct

contributor-book/src/guides/adding-a-new-operation-to-burn.md

@@ -79,11 +79,12 @@ previous section). Tests for `q_*` ops follow a similar procedure: the test is a
 the module name is inserted into `crates/burn-tensor/src/tests/quantization/ops/mod.rs` and finally
 the test is added to the
 [`testgen_quantization` macro](https://github.com/tracel-ai/burn/blob/a6a5c22e0db56d947b9165d4dae42783a5a6b689/crates/burn-tensor/src/tests/mod.rs#L67).
-If you take a look at any of the existing tests for an operation on a quantized tensor (e.g.,
-[`q_add`](https://github.com/tracel-ai/burn/blob/a6a5c22e0db56d947b9165d4dae42783a5a6b689/crates/burn-tensor/src/tests/quantization/ops/add.rs)),
+If you take a look at any of the existing tests for an operation on a quantized tensor,
 you will see that the inputs and expected outputs are always defined with floating point values.
 While it assumes that the quantization and dequantization are correct, it makes the tests much more
-readable and easier to understand w.r.t. what is being tested.
+readable and easier to understand w.r.t. what is being tested. Effectively, the tests are there to
+ensure that a tensor operation is invariant to quantization (up to some quantization error, of
+course).
 
 _Note: the tests try to use tensors with floating point values which can be de/quantized without
 introducing too much quantization error, but the result always depends on the operation (e.g.,

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

@@ -66,6 +66,42 @@ macro_rules! testgen_all {
 #[macro_export]
 macro_rules! testgen_quantization {
     () => {
+        // Quantized tensor utilities
+        pub mod qtensor {
+            use core::marker::PhantomData;
+
+            use burn_tensor::{
+                backend::Backend,
+                quantization::{QuantizationScheme, QuantizationType},
+                Tensor, TensorData,
+            };
+
+            pub struct QTensor<B: Backend, const D: usize> {
+                b: PhantomData<B>,
+            }
+
+            impl<B: Backend, const D: usize> QTensor<B, D> {
+                /// Creates a quantized int8 tensor from the floating point data using the default quantization scheme
+                /// (i.e., per-tensor symmetric quantization).
+                pub fn int8<F: Into<TensorData>>(floats: F) -> Tensor<B, D> {
+                    Self::int8_symmetric(floats)
+                }
+                /// Creates a quantized int8 tensor from the floating point data using per-tensor symmetric quantization.
+                pub fn int8_symmetric<F: Into<TensorData>>(floats: F) -> Tensor<B, D> {
+                    Tensor::from_floats(floats, &Default::default()).quantize_dynamic(
+                        &QuantizationScheme::PerTensorSymmetric(QuantizationType::QInt8),
+                    )
+                }
+                /// Creates a quantized int8 tensor from the floating point data using per-tensor affine quantization.
+                pub fn int8_affine<F: Into<TensorData>>(floats: F) -> Tensor<B, D> {
+                    Tensor::from_floats(floats, &Default::default()).quantize_dynamic(
+                        &QuantizationScheme::PerTensorAffine(QuantizationType::QInt8),
+                    )
+                }
+            }
+        }
+        pub use qtensor::*;
+
         // test quantization
         burn_tensor::testgen_calibration!();
         burn_tensor::testgen_scheme!();
@@ -104,6 +140,7 @@ macro_rules! testgen_quantization {
         burn_tensor::testgen_q_remainder!();
         burn_tensor::testgen_q_repeat_dim!();
         burn_tensor::testgen_q_reshape!();
+        burn_tensor::testgen_q_round!();
         burn_tensor::testgen_q_select!();
         burn_tensor::testgen_q_sin!();
         burn_tensor::testgen_q_slice!();

crates/burn-tensor/src/tests/quantization/ops/abs.rs

@@ -1,18 +1,11 @@
 #[burn_tensor_testgen::testgen(q_abs)]
 mod tests {
     use super::*;
-    use burn_tensor::quantization::{QuantizationStrategy, SymmetricQuantization};
-    use burn_tensor::{Tensor, TensorData};
+    use burn_tensor::TensorData;
 
     #[test]
     fn should_support_abs_ops() {
-        // Quantized [[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]]
-        let data = TensorData::quantized(
-            vec![0i8, -25, 51, 76, 102, -127],
-            [2, 3],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.03937008)),
-        );
-        let tensor = TestTensor::<2>::from_data(data, &Default::default());
+        let tensor = QTensor::<TestBackend, 2>::int8([[0.0, -1.0, 2.0], [3.0, 4.0, -5.0]]);
 
         let output = tensor.abs();
 

crates/burn-tensor/src/tests/quantization/ops/add.rs

@@ -1,25 +1,12 @@
 #[burn_tensor_testgen::testgen(q_add)]
 mod tests {
     use super::*;
-    use burn_tensor::quantization::{QuantizationStrategy, SymmetricQuantization};
-    use burn_tensor::{Tensor, TensorData};
+    use burn_tensor::TensorData;
 
     #[test]
     fn test_add_d2() {
-        // Quantized [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
-        let data = TensorData::quantized(
-            vec![0i8, 25, 51, 76, 102, 127],
-            [2, 3],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.03937008)),
-        );
-        let tensor_1 = TestTensor::<2>::from_data(data, &Default::default());
-        // Quantized [[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]
-        let data = TensorData::quantized(
-            vec![69i8, 81, 92, 104, 115, 127],
-            [2, 3],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.08661418)),
-        );
-        let tensor_2 = TestTensor::<2>::from_data(data, &Default::default());
+        let tensor_1 = QTensor::<TestBackend, 2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
+        let tensor_2 = QTensor::<TestBackend, 2>::int8([[6.0, 7.0, 8.0], [9.0, 10.0, 11.0]]);
 
         let output = tensor_1 + tensor_2;
 
@@ -32,20 +19,8 @@ mod tests {
 
     #[test]
     fn test_add_broadcast() {
-        // Quantized [[0.0, 1.0, 2.0]]
-        let data = TensorData::quantized(
-            vec![0i8, 64, 127],
-            [1, 3],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.015748031)),
-        );
-        let tensor_1 = TestTensor::<2>::from_data(data, &Default::default());
-        // Quantized [[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]
-        let data = TensorData::quantized(
-            vec![48i8, 64, 79, 95, 111, 127],
-            [2, 3],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.062992126)),
-        );
-        let tensor_2 = TestTensor::<2>::from_data(data, &Default::default());
+        let tensor_1 = QTensor::<TestBackend, 2>::int8([[0.0, 1.0, 2.0]]);
+        let tensor_2 = QTensor::<TestBackend, 2>::int8([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0]]);
 
         let output = tensor_1 + tensor_2;
 
@@ -58,22 +33,10 @@ mod tests {
 
     #[test]
     fn test_add_different_strides_rhs() {
-        // Quantized [[0.0, 1.0], [2.0, 3.0]]
-        let data = TensorData::quantized(
-            vec![0i8, 42, 85, 127],
-            [2, 2],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.023622047)),
-        );
         // We need to execute an operation after `from data` to trigger inplace in some backends.
         // Which is the operation that might be problematic in this case.
-        let tensor_1 = TestTensor::<2>::from_data(data, &Default::default()) * 1;
-        // Quantized [[4.0, 5.0], [6.0, 7.0]]
-        let data = TensorData::quantized(
-            vec![73i8, 91, 109, 127],
-            [2, 2],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.05511811)),
-        );
-        let tensor_2 = TestTensor::<2>::from_data(data, &Default::default()) * 1;
+        let tensor_1 = QTensor::<TestBackend, 2>::int8([[0.0, 1.0], [2.0, 3.0]]) * 1;
+        let tensor_2 = QTensor::<TestBackend, 2>::int8([[4.0, 5.0], [6.0, 7.0]]) * 1;
 
         let output = tensor_1 + tensor_2.transpose();
 
@@ -86,22 +49,10 @@ mod tests {
 
     #[test]
     fn test_add_different_strides_lhs() {
-        // Quantized [[0.0, 1.0], [2.0, 3.0]]
-        let data = TensorData::quantized(
-            vec![0i8, 42, 85, 127],
-            [2, 2],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.023622047)),
-        );
         // We need to execute an operation after `from data` to trigger inplace in some backends.
         // Which is the operation that might be problematic in this case.
-        let tensor_1 = TestTensor::<2>::from_data(data, &Default::default()) * 1;
-        // Quantized [[4.0, 5.0], [6.0, 7.0]]
-        let data = TensorData::quantized(
-            vec![73i8, 91, 109, 127],
-            [2, 2],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.05511811)),
-        );
-        let tensor_2 = TestTensor::<2>::from_data(data, &Default::default()) * 1;
+        let tensor_1 = QTensor::<TestBackend, 2>::int8([[0.0, 1.0], [2.0, 3.0]]) * 1;
+        let tensor_2 = QTensor::<TestBackend, 2>::int8([[4.0, 5.0], [6.0, 7.0]]) * 1;
 
         let output = tensor_1.transpose() + tensor_2;
 
@@ -114,22 +65,10 @@ mod tests {
 
     #[test]
     fn test_add_different_strides_broadcast() {
-        // Quantized [[0.0, 1.0], [2.0, 3.0]]
-        let data = TensorData::quantized(
-            vec![0i8, 42, 85, 127],
-            [2, 2],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.023622047)),
-        );
         // We need to execute an operation after `from data` to trigger inplace in some backends.
         // Which is the operation that might be problematic in this case.
-        let tensor_1 = TestTensor::<2>::from_data(data, &Default::default()) * 1;
-        // Quantized [[4.0, 5.0]]
-        let data = TensorData::quantized(
-            vec![102i8, 127],
-            [1, 2],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.03937008)),
-        );
-        let tensor_2 = TestTensor::<2>::from_data(data, &Default::default()) * 1;
+        let tensor_1 = QTensor::<TestBackend, 2>::int8([[0.0, 1.0], [2.0, 3.0]]) * 1;
+        let tensor_2 = QTensor::<TestBackend, 2>::int8([[4.0, 5.0]]) * 1;
 
         let output = tensor_1.transpose() + tensor_2;
 
@@ -143,13 +82,7 @@ mod tests {
     #[test]
     fn should_support_add_scalar_ops() {
         let scalar = 2.0;
-        // Quantized [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]
-        let data = TensorData::quantized(
-            vec![0i8, 25, 51, 76, 102, 127],
-            [2, 3],
-            QuantizationStrategy::PerTensorSymmetricInt8(SymmetricQuantization::init(0.03937008)),
-        );
-        let tensor = TestTensor::<2>::from_data(data, &Default::default());
+        let tensor = QTensor::<TestBackend, 2>::int8([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]]);
 
         let output = tensor + scalar;