Add QTensorOps require_grad methods to avoid dequantizing

tracel-ai · laggui · Jul 15, 2024 · Jul 9, 2024 · Jul 9, 2024 · Jul 11, 2024
commit feea531751750ca098a78b9a2319026301d5a7f6
diff --git a/crates/burn-tensor/src/tensor/api/float.rs b/crates/burn-tensor/src/tensor/api/float.rs
@@ -271,9 +271,9 @@ where
         match &self.primitive {
             TensorPrimitive::Float(tensor) => B::float_is_require_grad(tensor),
             TensorPrimitive::QFloat {
-                tensor: _,
+                tensor,
                 strategy: _,
-            } => B::float_is_require_grad(&self.primitive.clone().tensor()),
+            } => B::q_is_require_grad(tensor),
         }
     }
 
@@ -282,10 +282,16 @@ where
     ///
     /// This function does nothing when autodiff is not enabled.
     pub fn set_require_grad(self, require_grad: bool) -> Self {
-        Self::new(TensorPrimitive::Float(B::float_set_require_grad(
-            self.primitive.tensor(),
-            require_grad,
-        )))
+        let primitive = match self.primitive {
+            TensorPrimitive::Float(tensor) => {
+                TensorPrimitive::Float(B::float_set_require_grad(tensor, require_grad))
+            }
+            TensorPrimitive::QFloat { tensor, strategy } => TensorPrimitive::QFloat {
+                tensor: B::q_set_require_grad(tensor, require_grad),
+                strategy,
+            },
+        };
+        Self::new(primitive)
     }
 
     /// Applies the relu function to the tensor.

diff --git a/crates/burn-tensor/src/tensor/ops/qtensor.rs b/crates/burn-tensor/src/tensor/ops/qtensor.rs
@@ -81,4 +81,19 @@ pub trait QTensorOps<B: Backend> {
         tensor: QuantizedTensor<B, D>,
         strategy: QuantizationStrategy,
     ) -> impl Future<Output = TensorData> + Send;
+
+    /// Sets the `require_grad` flag of a tensor.
+    fn q_set_require_grad<const D: usize>(
+        tensor: QuantizedTensor<B, D>,
+        _require_grad: bool,
+    ) -> QuantizedTensor<B, D> {
+        // Should only be overridden by autodiff backends.
+        tensor
+    }
+
+    /// Returns the `require_grad` flag of a tensor.
+    fn q_is_require_grad<const D: usize>(_tensor: &QuantizedTensor<B, D>) -> bool {
+        // Should only be overridden by autodiff backends.
+        false
+    }
 }