[QwixProvider] Remove ptq back-compat aliases; migrate callsites to boxed_param

qwix/_src/providers/ptq.py

@@ -24,13 +24,14 @@
 from qwix._src.utils import flax_util
 
 # Backwards-compatible aliases for existing ptq.* users.
+# TODO(guoren): safely remove after users move to boxed_param.
 BoxedParamProvider = boxed_param.BoxedParamProvider
 WithAux = boxed_param.WithAux
 create_quantized_param = boxed_param.create_quantized_param
 quantize_act = boxed_param.quantize_act
 
 
-class PtqProvider(BoxedParamProvider):
+class PtqProvider(boxed_param.BoxedParamProvider):
   """Quantization provider for PTQ.
 
   In PTQ mode, weights needs to be pre-quantized. However, Qwix doesn't know
@@ -94,7 +95,7 @@ def quantize_params(
       if allow_extra_params:
         continue
       raise ValueError(f'{path} is not found in the abstract_quantized_params.')
-    if isinstance(abs_param, WithAux):
+    if isinstance(abs_param, boxed_param.WithAux):
       # The param might not be in the shape needed for compute, in case the
       # module reshapes before compute. Abstract param has the compute shape.
       param = param.reshape(abs_param.shape)
@@ -112,12 +113,13 @@ def quantize_params(
     if quant_stat['count'] == 0:
       raise ValueError(f'quant_stats is not initialized for {path}.')
 
-    # Get the act_qtype from the scale, which is a WithAux[jax.Array].
+    # Get the act_qtype from the scale, which is a
+    # boxed_param.WithAux[jax.Array].
     scale_path = (*path[:-1], path[-1] + '_scale')
     abs_scale = flax_util.get_value_from_path(
         abstract_quantized_params, scale_path
     )
-    assert isinstance(abs_scale, WithAux)
+    assert isinstance(abs_scale, boxed_param.WithAux)
     act_qtype = abs_scale.how.qtype
 
     calibration = averaging.SimpleMovingAverage().get_calibration(quant_stat)
@@ -129,7 +131,7 @@ def quantize_params(
       quantized_params[(*path[:-1], path[-1] + '_zero_point')] = zero_point
 
   if isinstance(abstract_quantized_params, nnx.Module):
-    # Convert WithAux to a pure dict so that nnx.update() can work.
+    # Convert boxed_param.WithAux to a pure dict so that nnx.update() can work.
     quantized_params = nnx.to_pure_dict(nnx.state(quantized_params))
 
   return flax.traverse_util.unflatten_dict(quantized_params)

qwix/_src/utils/flax_util.py

@@ -206,12 +206,12 @@ def find_param(x: Any, ptq_array_type=None) -> str | None:
   equally.
 
   Args:
-    x: jax.Array or array-like object (such as ptq.WithAux) that has the "shape"
-      attribute.
+    x: jax.Array or array-like object (such as boxed_param.WithAux) that has the
+      "shape" attribute.
     ptq_array_type: when looking for a param in the current module, also
       consider this type in addition to jax.Array. This is only used in QLoRA
-      mode and should be set to ptq.WithAux. We don't import ptq here to avoid a
-      circular dependency.
+      mode and should be set to boxed_param.WithAux. We don't import ptq here to
+      avoid a circular dependency.
 
   Returns:
     The name of the param that contains the given array, or None if not found.
@@ -347,7 +347,7 @@ def update_sharding(
     spec = sum(spec, ())  # flatten the list of tuples.
   elif merge:
     for i in merge:
-      spec = spec[: i + 1] + spec[i + 2 :]  # pytype: disable=unsupported-operands # pyrefly: ignore
+      spec = spec[: i + 1] + spec[i + 2 :]  # pyrefly: ignore[unsupported-operation]
   elif transpose:
     spec = tuple(spec[i] if i is not None else None for i in transpose)
 

qwix/contrib/awq.py

@@ -37,6 +37,7 @@
 import jax
 from qwix._src import qconfig
 from qwix._src.core import qarray
+from qwix._src.providers import boxed_param
 from qwix._src.providers import ptq
 from qwix.contrib import awq_core
 from qwix.contrib import calibration
@@ -56,7 +57,7 @@ class AwqRule(qconfig.QuantizationRule):
 
 
 @flax.struct.dataclass(kw_only=True)
-class WithAwqScale(ptq.WithAux[qarray.QArray]):
+class WithAwqScale(boxed_param.WithAux[qarray.QArray]):
   """A quantized array with AWQ per-channel scales.
 
   This wrapper stores the quantized weights along with the per-channel AWQ
@@ -72,6 +73,7 @@ class WithAwqScale(ptq.WithAux[qarray.QArray]):
   awq_scale: jax.Array
   contracting_axis: int = flax.struct.field(pytree_node=False)
 
+
 # Register as NNX data to allow JAX arrays in Module attributes.
 nnx.register_data_type(WithAwqScale)
 
@@ -124,6 +126,7 @@ def quantize_params(
     weights, returns WithAwqScale wrappers containing the QArray and per-channel
     AWQ scales. For non-AWQ weights, returns WithAux wrappers (same as PTQ).
   """
+
   def _quantize(ctx: calibration.CalibratedQuantContext) -> Any:
     activation_scale = ctx.calibration_stats['act_scale']
     assert activation_scale.shape[0] == ctx.weight.shape[1]
@@ -182,7 +185,7 @@ def _apply_awq_scale(self, rhs: WithAwqScale) -> jax.Array:
   def dot_general(
       self,
       lhs: jax.Array,
-      rhs: jax.Array | WithAwqScale | ptq.WithAux[qarray.QArray],
+      rhs: jax.Array | WithAwqScale | boxed_param.WithAux[qarray.QArray],
       dimension_numbers: jax.lax.DotDimensionNumbers,
       precision: jax.lax.PrecisionLike = None,
       preferred_element_type: jax.typing.DTypeLike | None = None,
@@ -220,7 +223,7 @@ def _preprocess_operand(op):
     new_operands = jax.tree.map(
         _preprocess_operand,
         operands,
-        is_leaf=lambda x: isinstance(x, (WithAwqScale, ptq.WithAux)),
+        is_leaf=lambda x: isinstance(x, (WithAwqScale, boxed_param.WithAux)),
     )
 
     return super().einsum(

qwix/contrib/calibration.py

@@ -23,6 +23,7 @@
 from jax import numpy as jnp
 from qwix._src import averaging
 from qwix._src import qconfig
+from qwix._src.providers import boxed_param
 from qwix._src.providers import ptq
 from qwix._src.utils import flax_util
 
@@ -247,7 +248,7 @@ class CalibratedQuantContext:
   weight: jax.Array
   how: Any
   calibration_stats: dict[str, jax.Array]
-  abs_w: ptq.WithAux
+  abs_w: boxed_param.WithAux
   contracting_axis: int
   restore_shape: Callable[..., Any]
   path: tuple[str, ...]
@@ -256,7 +257,7 @@ class CalibratedQuantContext:
 def extract_calibrated_quant_context(
     path: tuple[str, ...],
     weight: jax.Array,
-    abs_w: ptq.WithAux,
+    abs_w: boxed_param.WithAux,
     stats: Any,
 ) -> CalibratedQuantContext | None:
   """Extracts the calibration context for a single weight.
@@ -338,7 +339,7 @@ def quantize_params_with_calibration(
     stats_path = (*path[:-1], path[-1] + stats_suffix)
     stats = flax_util.get_value_from_path(quant_stats, stats_path)
 
-    if not isinstance(abs_w, ptq.WithAux) or stats is None:
+    if not isinstance(abs_w, boxed_param.WithAux) or stats is None:
       not_quantized_params[path] = w
       continue
 

qwix/contrib/padded_ptq.py

@@ -33,6 +33,7 @@
 from qwix._src.core import einsum as core_einsum
 from qwix._src.core import einsum_info as core_einsum_info
 from qwix._src.core import qarray
+from qwix._src.providers import boxed_param as _boxed_param
 from qwix._src.providers import ptq as _ptq
 
 
@@ -276,7 +277,7 @@ def quantize_act(
     act_name: str | None,
 ):
   """Wrapper to reuse PTQ.quantize_act with this module as qarray backend."""
-  return _ptq.quantize_act(
+  return _boxed_param.quantize_act(
       array, how, rule, act_name, _qarray_module=sys.modules[__name__]  # pyrefly: ignore[bad-argument-type]
   )
 
@@ -285,9 +286,9 @@ def create_quantized_param(
     name: str,
     value: jax.Array,
     how: HowToQuantize,
-) -> _ptq.WithAux[qarray.QArray]:
+) -> _boxed_param.WithAux[qarray.QArray]:
   """Wrapper that delegates to PTQ.create_quantized_param using this backend."""
-  return _ptq.create_quantized_param(
+  return _boxed_param.create_quantized_param(
       name, value, how, _qarray_module=sys.modules[__name__]
   )
 

qwix/contrib/qep.py

@@ -37,6 +37,7 @@
 from qwix._src import model as qwix_model
 from qwix._src import qconfig
 from qwix._src.core import qarray
+from qwix._src.providers import boxed_param
 from qwix._src.providers import ptq
 from qwix.contrib import calibration
 from qwix.contrib import gptq
@@ -633,7 +634,7 @@ def _quantize_weight(
     rule: QepRule,
     gptq_block_size: int,
     gptq_damping_factor: float,
-) -> ptq.WithAux:
+) -> boxed_param.WithAux:
   """Generates a compressed discrete weight, adapting the QEP formula.
 
   Extracts the raw floating weight, adjusts it using the QEP compensation metric

qwix/contrib/smooth_quant.py

@@ -31,6 +31,7 @@
 from qwix._src import qconfig
 from qwix._src.core import dot_general
 from qwix._src.core import qarray
+from qwix._src.providers import boxed_param
 from qwix._src.providers import ptq
 from qwix._src.utils import flax_util
 from qwix.contrib import calibration
@@ -57,7 +58,7 @@ class SqRule(qconfig.QuantizationRule):
 
 
 @flax.struct.dataclass(kw_only=True)
-class WithSqScale(ptq.WithAux[qarray.QArray]):
+class WithSqScale(boxed_param.WithAux[qarray.QArray]):
   """A quantized array with SQ per-channel scales.
 
   This wrapper stores the quantized weights along with the per-channel SQ
@@ -190,7 +191,7 @@ def quantize_params(
     sq_stats_path = (*path[:-1], path[-1] + "_sq")
     sq_stats = flax_util.get_value_from_path(sq_quant_stats, sq_stats_path)
 
-    if not isinstance(abs_w, ptq.WithAux) or sq_stats is None:
+    if not isinstance(abs_w, boxed_param.WithAux) or sq_stats is None:
       # Not quantized by SQ.
       not_quantized_params[path] = w
       continue
@@ -264,7 +265,7 @@ def _apply_sq_scale(
   def dot_general(
       self,
       lhs: jax.Array,
-      rhs: jax.Array | WithSqScale | ptq.WithAux[qarray.QArray],
+      rhs: jax.Array | WithSqScale | boxed_param.WithAux[qarray.QArray],
       dimension_numbers: jax.lax.DotDimensionNumbers,
       precision: jax.lax.PrecisionLike = None,
       preferred_element_type: jax.typing.DTypeLike | None = None,

tests/_src/providers/lora_test.py

@@ -21,8 +21,8 @@
 import jax
 from jax import numpy as jnp
 from jax import sharding as shd
+from qwix._src.providers import boxed_param
 from qwix._src.providers import lora
-from qwix._src.providers import ptq
 
 os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4"
 
@@ -342,7 +342,7 @@ def test_lora_dot_general_nn(self, weight_qtype):
       self.assertEqual(kernel.unbox().shape, (16, 32))
       self.assertEqual(kernel.names, ("a", "b"))
     else:
-      self.assertIsInstance(kernel, ptq.WithAux)
+      self.assertIsInstance(kernel, boxed_param.WithAux)
       self.assertIsInstance(kernel.array.qvalue, nn.Partitioned)
       self.assertEqual(kernel.array.qvalue.unbox().shape, (16, 32))
       self.assertEqual(kernel.array.qvalue.names, ("a", "b"))