From c1e03b311cde99b12812c8508f336e9f5da8e222 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Thu, 18 Dec 2025 23:58:09 +0100
Subject: [PATCH 01/11] qq linear

---
 python/mlx/nn/layers/__init__.py  |  7 ++-
 python/mlx/nn/layers/quantized.py | 86 +++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 1 deletion(-)

diff --git a/python/mlx/nn/layers/__init__.py b/python/mlx/nn/layers/__init__.py
index ea2d3029d8..c2fba58347 100644
--- a/python/mlx/nn/layers/__init__.py
+++ b/python/mlx/nn/layers/__init__.py
@@ -87,7 +87,12 @@
     MaxPool3d,
 )
 from mlx.nn.layers.positional_encoding import ALiBi, RoPE, SinusoidalPositionalEncoding
-from mlx.nn.layers.quantized import QuantizedEmbedding, QuantizedLinear, quantize
+from mlx.nn.layers.quantized import (
+    QQLinear,
+    QuantizedEmbedding,
+    QuantizedLinear,
+    quantize,
+)
 from mlx.nn.layers.recurrent import GRU, LSTM, RNN
 from mlx.nn.layers.transformer import (
     MultiHeadAttention,
diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index c308e884ba..73056c10c9 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -268,3 +268,89 @@ def from_linear(
             ql.bias = linear_layer.bias
 
         return ql
+
+
+class QQLinear(Module):
+    """Quantizes input and applies an affine transformation to it
+        using a quantized weight matrix.
+
+    Weights can be either in quantized form (together with itss scales)
+    or in higher precision. If the weights are expected to be included
+    in gradient computations they are stored in higher precision
+    and quantized on the fly during computation.
+
+    Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
+    quantizes the input as well and includes weights in gradient computations.
+
+    :obj:`QQLinear` also provides the class method :meth:`from_linear` to convert
+    :class:`mlx.nn.Linear` layers to :obj:`QQLinear` layers.
+
+    Note: This layer does not support bias terms yet.
+    Args:
+        input_dims (int): The dimensionality of the input features.
+        output_dims (int): The dimensionality of the output features.
+        group_size (int, optional): The group size to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``16``.
+        bits (int, optional): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``4``.
+        mode (str, optional): The quantization method to use (see
+            :func:`mlx.core.quantize`). Currently, only ``"nvfp4"`` and ``"mxfp8"``
+            are supported. Default: ``"nvfp4"``.
+    """
+
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        group_size: int = 16,
+        bits: int = 4,
+        mode: str = "nvfp4",
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size = group_size
+        self.bits = bits
+        self.mode = mode
+
+        scale = math.sqrt(1 / input_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(output_dims, input_dims),
+        )
+
+    def _extra_repr(self):
+        out_dims, in_dims = self.weight.shape
+        if self.weight.dtype == mx.uint32:
+            in_dims *= 32 // self.bits
+        return (
+            f"input_dims={in_dims}, output_dims={out_dims}, "
+            f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
+        )
+
+    def __call__(self, x):
+        x = mx.qqmm(
+            x,
+            self["weight"],
+            getattr(self, "scales", None),
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )
+        return x
+
+    @classmethod
+    def from_linear(
+        cls,
+        linear_layer: Module,
+        group_size: int = 16,
+        bits: int = 4,
+        mode: str = "nvfp4",
+    ):
+        """Create a :obj:`QQLinear` layer from a :obj:`Linear` layer."""
+        output_dims, input_dims = linear_layer.weight.shape  # (N,K)
+        ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
+        ql.weight = linear_layer.weight
+
+        return ql

From 867e0fc45e4f581da17ecd0d65723db894c88ca5 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 19 Dec 2025 00:06:51 +0100
Subject: [PATCH 02/11] qq linear

---
 python/mlx/nn/layers/quantized.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 73056c10c9..79ffed01b5 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -274,10 +274,10 @@ class QQLinear(Module):
     """Quantizes input and applies an affine transformation to it
         using a quantized weight matrix.
 
-    Weights can be either in quantized form (together with itss scales)
-    or in higher precision. If the weights are expected to be included
-    in gradient computations they are stored in higher precision
-    and quantized on the fly during computation.
+    The weight matrix can be stored either in quantized form together 
+    with its scales. If the weights are expected to be included in 
+    gradient computations they are stored in higher precision and 
+    quantized on the fly during computation.
 
     Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
     quantizes the input as well and includes weights in gradient computations.

From 592da25375e115ac6c5f582635b84d661e7efa55 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 19 Dec 2025 00:07:08 +0100
Subject: [PATCH 03/11] Revert "qq linear"

This reverts commit 867e0fc45e4f581da17ecd0d65723db894c88ca5.
---
 python/mlx/nn/layers/quantized.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 79ffed01b5..73056c10c9 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -274,10 +274,10 @@ class QQLinear(Module):
     """Quantizes input and applies an affine transformation to it
         using a quantized weight matrix.
 
-    The weight matrix can be stored either in quantized form together 
-    with its scales. If the weights are expected to be included in 
-    gradient computations they are stored in higher precision and 
-    quantized on the fly during computation.
+    Weights can be either in quantized form (together with itss scales)
+    or in higher precision. If the weights are expected to be included
+    in gradient computations they are stored in higher precision
+    and quantized on the fly during computation.
 
     Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
     quantizes the input as well and includes weights in gradient computations.

From 80b9496da7c01d9294627aa2c34f4bbcd70d895d Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 19 Dec 2025 00:18:59 +0100
Subject: [PATCH 04/11] corrected the docs

---
 python/mlx/nn/layers/quantized.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 73056c10c9..44481080be 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -271,13 +271,16 @@ def from_linear(
 
 
 class QQLinear(Module):
-    """Quantizes input and applies an affine transformation to it
-        using a quantized weight matrix.
+    """Quantizes the input and applies an affine transformation using quantized weights.
 
-    Weights can be either in quantized form (together with itss scales)
-    or in higher precision. If the weights are expected to be included
-    in gradient computations they are stored in higher precision
-    and quantized on the fly during computation.
+    Two use cases are supported:
+
+    1) **Inference / frozen weights**: weights are stored in quantized form together with
+       their scales (``self.weight`` is quantized and ``self.scales`` is provided).
+
+    2) **Training / weights are included in gradient computation**:
+        weights are stored in higher precision and are quantized on
+        the fly during computation.
 
     Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
     quantizes the input as well and includes weights in gradient computations.
@@ -285,7 +288,8 @@ class QQLinear(Module):
     :obj:`QQLinear` also provides the class method :meth:`from_linear` to convert
     :class:`mlx.nn.Linear` layers to :obj:`QQLinear` layers.
 
-    Note: This layer does not support bias terms yet.
+    Note: This layer does not support a bias term yet.
+
     Args:
         input_dims (int): The dimensionality of the input features.
         output_dims (int): The dimensionality of the output features.

From aebd5219d9d0ee10ca3ba31d0ad2a96c70aa0236 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 26 Dec 2025 17:40:20 +0100
Subject: [PATCH 05/11] set default to None, added eval and train methods

---
 python/mlx/nn/layers/quantized.py | 51 +++++++++++++++++++++++--------
 1 file changed, 39 insertions(+), 12 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 5b64b13a6b..75e39bbb8c 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -294,36 +294,39 @@ class QQLinear(Module):
     Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
     quantizes the input as well and includes weights in gradient computations.
 
-    :obj:`QQLinear` also provides the class method :meth:`from_linear` to convert
-    :class:`mlx.nn.Linear` layers to :obj:`QQLinear` layers.
+    :obj:`QQLinear` also provides:
+     -  the class method :meth:`from_linear` to convert :class:`mlx.nn.Linear` 
+     layers to :obj:`QQLinear` layers. If the layer is created in evaluation mode, 
+     the weights will be quantized.
+     -  the methods :meth:`eval` and :meth:`train` to switch between inference
+        and training modes.
 
     Note: This layer does not support a bias term yet.
 
     Args:
         input_dims (int): The dimensionality of the input features.
         output_dims (int): The dimensionality of the output features.
-        group_size (int, optional): The group size to use for the quantized weight.
-            See :func:`~mlx.core.quantize`. Default: ``16``.
-        bits (int, optional): The bit width to use for the quantized weight.
-            See :func:`~mlx.core.quantize`. Default: ``4``.
-        mode (str, optional): The quantization method to use (see
+        group_size (Optional[int]): The group size to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        bits (Optional[int]): The bit width to use for the quantized weight.
+            See :func:`~mlx.core.quantize`. Default: ``None``.
+        mode (Optional[str]): The quantization method to use (see
             :func:`mlx.core.quantize`). Currently, only ``"nvfp4"`` and ``"mxfp8"``
-            are supported. Default: ``"nvfp4"``.
+            are supported. Default: ``None``.
     """
 
     def __init__(
         self,
         input_dims: int,
         output_dims: int,
-        group_size: int = 16,
-        bits: int = 4,
+        group_size: int = None,
+        bits: int = None,
         mode: str = "nvfp4",
     ):
         super().__init__()
 
         # Quantization config
-        self.group_size = group_size
-        self.bits = bits
+        self.group_size, self.bits = _defaults_for_mode(mode, group_size, bits)
         self.mode = mode
 
         scale = math.sqrt(1 / input_dims)
@@ -342,6 +345,26 @@ def _extra_repr(self):
             f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
         )
 
+    def eval(self):
+        if self.weight.dtype != mx.uint32:
+            self.weight, self.scales = mx.quantize(
+                self.weight,
+                self.group_size,
+                self.bits,
+                mode=self.mode,
+            )
+    
+    def train(self):
+        if self.weight.dtype == mx.uint32:
+            self.weight = mx.dequantize(
+                self.weight,
+                scales=self.scales,
+                group_size=self.group_size,
+                bits=self.bits,
+                mode=self.mode,
+            )
+            del self.scales
+
     def __call__(self, x):
         x = mx.qqmm(
             x,
@@ -360,10 +383,14 @@ def from_linear(
         group_size: int = 16,
         bits: int = 4,
         mode: str = "nvfp4",
+        train: bool = True,
     ):
         """Create a :obj:`QQLinear` layer from a :obj:`Linear` layer."""
         output_dims, input_dims = linear_layer.weight.shape  # (N,K)
         ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
         ql.weight = linear_layer.weight
 
+        if not train:
+            ql.eval()
+
         return ql

From 949b0f2c8af91000f83e23f14c46af8392608d20 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 26 Dec 2025 17:55:24 +0100
Subject: [PATCH 06/11] to_linear defaults to None

---
 python/mlx/nn/layers/quantized.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 75e39bbb8c..0a5d1cedf9 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -380,8 +380,8 @@ def __call__(self, x):
     def from_linear(
         cls,
         linear_layer: Module,
-        group_size: int = 16,
-        bits: int = 4,
+        group_size: int = None,
+        bits: int = None,
         mode: str = "nvfp4",
         train: bool = True,
     ):

From 056bc600ec0e584897e5f75f35135cbd696361b4 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 26 Dec 2025 18:09:15 +0100
Subject: [PATCH 07/11] typo

---
 python/mlx/nn/layers/quantized.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 0a5d1cedf9..4664789e05 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -312,7 +312,7 @@ class QQLinear(Module):
             See :func:`~mlx.core.quantize`. Default: ``None``.
         mode (Optional[str]): The quantization method to use (see
             :func:`mlx.core.quantize`). Currently, only ``"nvfp4"`` and ``"mxfp8"``
-            are supported. Default: ``None``.
+            are supported. Default: ``nvfp4``.
     """
 
     def __init__(

From fac92e07c29b44bfdcc21acc55e358b2fc6f1733 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 26 Dec 2025 18:09:46 +0100
Subject: [PATCH 08/11] re-commit

---
 python/mlx/nn/layers/quantized.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 4664789e05..cc58660b6d 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -295,8 +295,8 @@ class QQLinear(Module):
     quantizes the input as well and includes weights in gradient computations.
 
     :obj:`QQLinear` also provides:
-     -  the class method :meth:`from_linear` to convert :class:`mlx.nn.Linear` 
-     layers to :obj:`QQLinear` layers. If the layer is created in evaluation mode, 
+     -  the class method :meth:`from_linear` to convert :class:`mlx.nn.Linear`
+     layers to :obj:`QQLinear` layers. If the layer is created in evaluation mode,
      the weights will be quantized.
      -  the methods :meth:`eval` and :meth:`train` to switch between inference
         and training modes.
@@ -353,7 +353,7 @@ def eval(self):
                 self.bits,
                 mode=self.mode,
             )
-    
+
     def train(self):
         if self.weight.dtype == mx.uint32:
             self.weight = mx.dequantize(

From c2e522645ef6a470b2898ad62cc90f827ca23d86 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 2 Jan 2026 14:58:17 +0100
Subject: [PATCH 09/11] made api consistent with Module::train, docs changes

---
 python/mlx/nn/layers/quantized.py | 69 ++++++++++++++++---------------
 1 file changed, 35 insertions(+), 34 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index cc58660b6d..145630b75f 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -284,22 +284,20 @@ class QQLinear(Module):
 
     Two use cases are supported:
 
-    1) **Inference / frozen weights**: weights are stored in quantized form together with
+    1) **Eval**:  The weights are frozen and stored in quantized form together with
        their scales (``self.weight`` is quantized and ``self.scales`` is provided).
+    2) **Train**: The weights are stored in higher precision and are quantized on
+         the fly during computation so that gradients with respect to the weights
+         can be computed.
 
-    2) **Training / weights are included in gradient computation**:
-        weights are stored in higher precision and are quantized on
-        the fly during computation.
+    To switch between the two cases, use ``layer.eval()`` and ``layer.train()`` respectively.
 
     Compared to the :class:`mlx.nn.QuantizedLinear` layer, this layer
     quantizes the input as well and includes weights in gradient computations.
 
     :obj:`QQLinear` also provides:
      -  the class method :meth:`from_linear` to convert :class:`mlx.nn.Linear`
-     layers to :obj:`QQLinear` layers. If the layer is created in evaluation mode,
-     the weights will be quantized.
-     -  the methods :meth:`eval` and :meth:`train` to switch between inference
-        and training modes.
+     layers to :obj:`QQLinear` layers.
 
     Note: This layer does not support a bias term yet.
 
@@ -312,7 +310,7 @@ class QQLinear(Module):
             See :func:`~mlx.core.quantize`. Default: ``None``.
         mode (Optional[str]): The quantization method to use (see
             :func:`mlx.core.quantize`). Currently, only ``"nvfp4"`` and ``"mxfp8"``
-            are supported. Default: ``nvfp4``.
+            are supported. Default: ``"nvfp4"``.
     """
 
     def __init__(
@@ -335,6 +333,7 @@ def __init__(
             high=scale,
             shape=(output_dims, input_dims),
         )
+        self._quantized = False
 
     def _extra_repr(self):
         out_dims, in_dims = self.weight.shape
@@ -345,31 +344,34 @@ def _extra_repr(self):
             f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
         )
 
-    def eval(self):
-        if self.weight.dtype != mx.uint32:
-            self.weight, self.scales = mx.quantize(
-                self.weight,
-                self.group_size,
-                self.bits,
-                mode=self.mode,
-            )
-
-    def train(self):
-        if self.weight.dtype == mx.uint32:
-            self.weight = mx.dequantize(
-                self.weight,
-                scales=self.scales,
-                group_size=self.group_size,
-                bits=self.bits,
-                mode=self.mode,
-            )
-            del self.scales
+    def _update_mode(self):
+        if self._training:
+            if self._quantized:
+                self.weight = mx.dequantize(
+                    self.weight,
+                    scales=self.scales,
+                    group_size=self.group_size,
+                    bits=self.bits,
+                    mode=self.mode,
+                )
+                self.__delattr__("scales")
+                self._quantized = False
+        else:
+            if not self._quantized:
+                self.weight, self.scales = mx.quantize(
+                    self.weight,
+                    self.group_size,
+                    self.bits,
+                    mode=self.mode,
+                )
+                self._quantized = True
 
     def __call__(self, x):
+        self._update_mode()
         x = mx.qqmm(
             x,
             self["weight"],
-            getattr(self, "scales", None),
+            scales=self.get("scales"),
             group_size=self.group_size,
             bits=self.bits,
             mode=self.mode,
@@ -383,14 +385,13 @@ def from_linear(
         group_size: int = None,
         bits: int = None,
         mode: str = "nvfp4",
-        train: bool = True,
     ):
         """Create a :obj:`QQLinear` layer from a :obj:`Linear` layer."""
         output_dims, input_dims = linear_layer.weight.shape  # (N,K)
+        if linear_layer.get('bias') is not None:
+            raise NotImplementedError("QQLinear does not support bias yet.")
         ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
         ql.weight = linear_layer.weight
-
-        if not train:
-            ql.eval()
-
+        ql._training = linear_layer._training
+            
         return ql

From 394ab77175a467a8de1e7fdaf2d7899b10cc0d54 Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Fri, 2 Jan 2026 15:19:31 +0100
Subject: [PATCH 10/11] pre-commit

---
 python/mlx/nn/layers/quantized.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 145630b75f..54b59dd6eb 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -388,10 +388,10 @@ def from_linear(
     ):
         """Create a :obj:`QQLinear` layer from a :obj:`Linear` layer."""
         output_dims, input_dims = linear_layer.weight.shape  # (N,K)
-        if linear_layer.get('bias') is not None:
+        if linear_layer.get("bias") is not None:
             raise NotImplementedError("QQLinear does not support bias yet.")
         ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
         ql.weight = linear_layer.weight
         ql._training = linear_layer._training
-            
+
         return ql

From 86879cb8aef78a5ba699b1b4ac7b5ebd2fbd4c1a Mon Sep 17 00:00:00 2001
From: Anastasiia Filippova <a_filippova@apple.com>
Date: Mon, 5 Jan 2026 16:19:27 +0100
Subject: [PATCH 11/11] set training mode

---
 python/mlx/nn/layers/base.py      |  7 +++--
 python/mlx/nn/layers/quantized.py | 49 ++++++++++++++++++-------------
 2 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/python/mlx/nn/layers/base.py b/python/mlx/nn/layers/base.py
index ce009d9f83..fca65d787d 100644
--- a/python/mlx/nn/layers/base.py
+++ b/python/mlx/nn/layers/base.py
@@ -559,6 +559,9 @@ def _unfreeze_impl(_, m):
             _unfreeze_impl("", self)
         return self
 
+    def _set_training_mode(self, mode: bool) -> None:
+        self._training = mode
+
     def train(self, mode: bool = True) -> Module:
         """Set the model in or out of training mode.
 
@@ -573,10 +576,8 @@ def train(self, mode: bool = True) -> Module:
             The module instance after updating the training mode.
         """
 
-        def _set_train(_, m):
-            m._training = mode
+        self.apply_to_modules(lambda _, m: m._set_training_mode(mode))
 
-        self.apply_to_modules(_set_train)
         return self
 
     def eval(self) -> Module:
diff --git a/python/mlx/nn/layers/quantized.py b/python/mlx/nn/layers/quantized.py
index 54b59dd6eb..1c98706daf 100644
--- a/python/mlx/nn/layers/quantized.py
+++ b/python/mlx/nn/layers/quantized.py
@@ -344,30 +344,37 @@ def _extra_repr(self):
             f"group_size={self.group_size}, bits={self.bits}, mode={self.mode}"
         )
 
-    def _update_mode(self):
+    def quantize(self):
+        if not self._quantized:
+            self.weight, self.scales = mx.quantize(
+                self.weight,
+                self.group_size,
+                self.bits,
+                mode=self.mode,
+            )
+            self._quantized = True
+
+    def dequantize(self):
+        if self._quantized:
+            self.weight = mx.dequantize(
+                self.weight,
+                scales=self.scales,
+                group_size=self.group_size,
+                bits=self.bits,
+                mode=self.mode,
+            )
+            self.__delattr__("scales")
+            self._quantized = False
+
+    def _set_training_mode(self, mode: bool):
+        super()._set_training_mode(mode)
+
         if self._training:
-            if self._quantized:
-                self.weight = mx.dequantize(
-                    self.weight,
-                    scales=self.scales,
-                    group_size=self.group_size,
-                    bits=self.bits,
-                    mode=self.mode,
-                )
-                self.__delattr__("scales")
-                self._quantized = False
+            self.dequantize()
         else:
-            if not self._quantized:
-                self.weight, self.scales = mx.quantize(
-                    self.weight,
-                    self.group_size,
-                    self.bits,
-                    mode=self.mode,
-                )
-                self._quantized = True
+            self.quantize()
 
     def __call__(self, x):
-        self._update_mode()
         x = mx.qqmm(
             x,
             self["weight"],
@@ -392,6 +399,6 @@ def from_linear(
             raise NotImplementedError("QQLinear does not support bias yet.")
         ql = cls(input_dims, output_dims, group_size, bits, mode=mode)
         ql.weight = linear_layer.weight
-        ql._training = linear_layer._training
+        ql.train(linear_layer.training)
 
         return ql