diff --git a/qwix/_src/providers/ptq.py b/qwix/_src/providers/ptq.py
index 0ad807c..42be97c 100644
--- a/qwix/_src/providers/ptq.py
+++ b/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)
diff --git a/qwix/_src/utils/flax_util.py b/qwix/_src/utils/flax_util.py
index 78f605d..e680886 100644
--- a/qwix/_src/utils/flax_util.py
+++ b/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)
 
diff --git a/qwix/contrib/awq.py b/qwix/contrib/awq.py
index 0a78203..d79d898 100644
--- a/qwix/contrib/awq.py
+++ b/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(
diff --git a/qwix/contrib/calibration.py b/qwix/contrib/calibration.py
index 2a0dd49..b6298c7 100644
--- a/qwix/contrib/calibration.py
+++ b/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
 
diff --git a/qwix/contrib/padded_ptq.py b/qwix/contrib/padded_ptq.py
index b42579c..711e843 100644
--- a/qwix/contrib/padded_ptq.py
+++ b/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__]
   )
 
diff --git a/qwix/contrib/qep.py b/qwix/contrib/qep.py
index dd13b53..617a43e 100644
--- a/qwix/contrib/qep.py
+++ b/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
diff --git a/qwix/contrib/smooth_quant.py b/qwix/contrib/smooth_quant.py
index 4d6ba25..6c7ce86 100644
--- a/qwix/contrib/smooth_quant.py
+++ b/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,
diff --git a/tests/_src/providers/lora_test.py b/tests/_src/providers/lora_test.py
index 615a09b..8750ca1 100644
--- a/tests/_src/providers/lora_test.py
+++ b/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"))
diff --git a/tests/_src/providers/ptq_test.py b/tests/_src/providers/ptq_test.py
index 85303b9..2a5b629 100644
--- a/tests/_src/providers/ptq_test.py
+++ b/tests/_src/providers/ptq_test.py
@@ -24,6 +24,7 @@
 from jax.experimental import pallas as pl
 from qwix._src import model as qwix_model
 from qwix._src import qconfig
+from qwix._src.providers import boxed_param
 from qwix._src.providers import ptq
 from qwix._src.providers import qt
 from qwix._src.utils import flax_util
@@ -55,7 +56,7 @@ def test_nn_ptq(self):
     qw = ptq_abs_params["kernel"]
 
     # Test PTQ param structure.
-    self.assertIsInstance(qw, ptq.WithAux)
+    self.assertIsInstance(qw, boxed_param.WithAux)
     qw = qw.array
     self.assertIsInstance(qw.qvalue, nn.Partitioned)
     self.assertIsInstance(qw.scale, nn.Partitioned)
@@ -124,7 +125,9 @@ def test_nn_srq(self, act_calibration_method):
         ptq_dense.init, jax.random.key(0), model_input
     )["params"]
     self.assertIn("dot_general0_lhs_scale", ptq_abs_params)
-    self.assertIsInstance(ptq_abs_params["dot_general0_lhs_scale"], ptq.WithAux)
+    self.assertIsInstance(
+        ptq_abs_params["dot_general0_lhs_scale"], boxed_param.WithAux
+    )
     self.assertEqual(
         ptq_abs_params["dot_general0_lhs_scale"].array.shape, (1, 1)
     )
@@ -174,7 +177,7 @@ def test_nnx_ptq(self):
       )
     # Test PTQ param structure.
     qw = ptq_linear.kernel
-    self.assertIsInstance(qw, ptq.WithAux)
+    self.assertIsInstance(qw, boxed_param.WithAux)
     qw = qw.array
     self.assertEqual(qw.qvalue.dtype, jnp.int8)
     self.assertEqual(qw.qvalue.shape, (12, 6))
@@ -265,7 +268,7 @@ def get_canonical_pspec(x: jax.Array):
 
     # Test PTQ param structure.
     qw = ptq_einsum.kernel
-    self.assertIsInstance(qw, ptq.WithAux)
+    self.assertIsInstance(qw, boxed_param.WithAux)
     qw = qw.array
     self.assertEqual(qw.qvalue.dtype, jnp.int8)
     self.assertEqual(qw.qvalue.shape, (16, 8, 10))
@@ -343,7 +346,9 @@ def test_nnx_srq(self, act_calibration_method):
     ptq_linear = qwix_model.quantize_model(
         qt_linear, ptq.PtqProvider(q_rules), model_input
     )
-    self.assertIsInstance(ptq_linear.dot_general0_lhs_scale, ptq.WithAux)
+    self.assertIsInstance(
+        ptq_linear.dot_general0_lhs_scale, boxed_param.WithAux
+    )
     self.assertEqual(ptq_linear.dot_general0_lhs_scale.array.shape, (1, 1))
     if act_calibration_method == "minmax":
       self.assertEqual(ptq_linear.dot_general0_lhs_zero_point.shape, (1, 1))
@@ -401,7 +406,7 @@ def pallas_dot(x, y, out):
     ]
     ptq_model = qwix_model.quantize_model(model, ptq.PtqProvider(q_rules))
     variables = ptq_model.init(jax.random.key(0), jnp.ones((16, 32)))
-    self.assertIsInstance(variables["params"]["w1"], ptq.WithAux)
+    self.assertIsInstance(variables["params"]["w1"], boxed_param.WithAux)
     self.assertIsInstance(variables["params"]["w2"], jax.Array)
     ptq_model.apply(variables, jnp.ones((16, 32)))
 
@@ -452,7 +457,7 @@ def scan_fn(x: jax.Array, layer):
     ptq_model = qwix_model.quantize_model(
         model, ptq.PtqProvider(q_rules), model_input
     )
-    self.assertIsInstance(ptq_model.layers.kernel, ptq.WithAux)
+    self.assertIsInstance(ptq_model.layers.kernel, boxed_param.WithAux)
     self.assertIsInstance(ptq_model.layers.kernel.array, ptq.qarray.QArray)
     self.assertEqual(ptq_model.layers.kernel.array.shape, (2, 12, 12))
     self.assertEqual(ptq_model.layers.kernel.array.qtype, jnp.int8)
@@ -489,7 +494,7 @@ def __call__(self, x):
     )
     variables["params"] = ptq_params
 
-    self.assertIsInstance(variables["params"]["w"], ptq.WithAux)
+    self.assertIsInstance(variables["params"]["w"], boxed_param.WithAux)
     self.assertEqual(variables["params"]["w"].array.shape, (4, 5))
 
     # Test can apply.
@@ -561,7 +566,7 @@ def test_asarray_interception_with_aux(self):
     model = self._get_quantized_asarray_model()
 
     # Case 1: WithAux
-    wa = ptq.WithAux(
+    wa = boxed_param.WithAux(
         jnp.ones((2, 2)),
         qconfig.QuantizationRule(weight_qtype=jnp.int8).weight_qtype,
     )
@@ -569,18 +574,18 @@ def test_asarray_interception_with_aux(self):
     self.assertEqual(jax.tree.map(id, res_wa), jax.tree.map(id, wa))
 
     # Case 1.1: WithAux with non-array content
-    wa_list = ptq.WithAux(
+    wa_list = boxed_param.WithAux(
         [1.0, 2.0],
         qconfig.QuantizationRule(weight_qtype=jnp.int8).weight_qtype,
     )
     res_wa_list = model(wa_list)
     self.assertIsInstance(res_wa_list.array, jax.Array)
-    self.assertIsInstance(res_wa_list, ptq.WithAux)
+    self.assertIsInstance(res_wa_list, boxed_param.WithAux)
 
     # Case 1.2: WithAux with explicit dtype
     res_wa_f16 = model(wa, dtype=jnp.float16)
     self.assertEqual(res_wa_f16.dtype, jnp.float16)
-    self.assertIsInstance(res_wa_f16, ptq.WithAux)
+    self.assertIsInstance(res_wa_f16, boxed_param.WithAux)
 
     # Case 1.3: Grouped properties (shape, ndim, dtype)
     self.assertEqual(wa.shape, (2, 2))
@@ -588,7 +593,7 @@ def test_asarray_interception_with_aux(self):
     self.assertEqual(wa.dtype, jnp.float32)
 
     # Case 1.4: Metadata preservation in astype
-    wa_partitioned = ptq.WithAux(
+    wa_partitioned = boxed_param.WithAux(
         nn.Partitioned(jnp.ones((2, 2)), names=("a", "b")),
         qconfig.QuantizationRule(weight_qtype=jnp.int8).weight_qtype,
     )
@@ -654,23 +659,23 @@ def test_asarray_with_aux_qarray(self):
     """asarray should handle WithAux wrapping a QArray."""
     model = self._get_quantized_asarray_model()
     qa = ptq.qarray.QArray(jnp.ones((2, 2), jnp.int8), jnp.ones((1, 1)))
-    wa = ptq.WithAux(
+    wa = boxed_param.WithAux(
         qa, qconfig.QuantizationRule(weight_qtype=jnp.int8).weight_qtype
     )
     res = model(wa)
-    self.assertIsInstance(res, ptq.WithAux)
+    self.assertIsInstance(res, boxed_param.WithAux)
     self.assertIsInstance(res.array, ptq.qarray.QArray)
 
   def test_asarray_nnx_param_with_aux(self):
     """asarray should handle nnx.Param wrapping a WithAux."""
     model = self._get_quantized_asarray_model()
-    wa = ptq.WithAux(
+    wa = boxed_param.WithAux(
         jnp.ones((2, 2)),
         qconfig.QuantizationRule(weight_qtype=jnp.int8).weight_qtype,
     )
     param = nnx.Param(wa)
     res = model(param)
-    self.assertIsInstance(res, ptq.WithAux)
+    self.assertIsInstance(res, boxed_param.WithAux)
 
   def test_nnx_multi_head_attention(self):
     q_rules = [
diff --git a/tests/_src/qconfig_test.py b/tests/_src/qconfig_test.py
index 780e72f..e0f7689 100644
--- a/tests/_src/qconfig_test.py
+++ b/tests/_src/qconfig_test.py
@@ -132,7 +132,6 @@ def test_boxed_param_provider_hierarchy(self):
         issubclass(lora.LoraProvider, boxed_param.BoxedParamProvider)
     )
     self.assertFalse(issubclass(lora.LoraProvider, ptq.PtqProvider))
-    self.assertIs(ptq.WithAux, boxed_param.WithAux)
 
 
 if __name__ == "__main__":
diff --git a/tests/_src/utils/flax_util_test.py b/tests/_src/utils/flax_util_test.py
index 9000b0a..e0098b8 100644
--- a/tests/_src/utils/flax_util_test.py
+++ b/tests/_src/utils/flax_util_test.py
@@ -20,7 +20,7 @@
 from jax import numpy as jnp
 from jax.nn import initializers
 import numpy as np
-from qwix._src.providers import ptq
+from qwix._src.providers import boxed_param
 from qwix._src.utils import flax_util
 
 
@@ -298,7 +298,7 @@ def __call__(self):
         t.assertEqual(flax_util.find_param(w), "w")
         t.assertEqual(flax_util.find_param(w.astype(jnp.bfloat16)), "w")
         t.assertEqual(flax_util.find_param(w.reshape((2, 2, 5))), "w")
-        t.assertEqual(flax_util.find_param(ptq.WithAux(w, None)), "w")
+        t.assertEqual(flax_util.find_param(boxed_param.WithAux(w, None)), "w")
 
     model = MyModule()
     variables = jax.jit(model.init)(jax.random.key(0))
@@ -320,7 +320,7 @@ def __call__(self):
         t.assertEqual(flax_util.find_param(self.w.astype(jnp.bfloat16)), "w")
         t.assertEqual(flax_util.find_param(self.w.reshape((2, 2, 5))), "w")
         t.assertEqual(
-            flax_util.find_param(ptq.WithAux(self.w.value, None)), "w"
+            flax_util.find_param(boxed_param.WithAux(self.w.value, None)), "w"
         )
 
     nnx.jit(MyModule())()
diff --git a/tests/contrib/padded_ptq_test.py b/tests/contrib/padded_ptq_test.py
index efe4af9..ec87e60 100644
--- a/tests/contrib/padded_ptq_test.py
+++ b/tests/contrib/padded_ptq_test.py
@@ -20,6 +20,7 @@
 from qwix._src import qconfig
 from qwix._src.core import dot_general as core_dot
 from qwix._src.core import einsum as core_einsum
+from qwix._src.providers import boxed_param
 from qwix._src.providers import ptq
 from qwix.contrib import padded_ptq
 
@@ -62,7 +63,7 @@ def test_einsum_simple(self, weight_qtype, act_qtype, tile_size):
     )
 
     w_qarray_pad = padded_ptq.quantize_act(w, how, rule, None)
-    w_qarray_ptq = ptq.quantize_act(w_padded, how, rule, None)
+    w_qarray_ptq = boxed_param.quantize_act(w_padded, how, rule, None)
 
     q_diff = jnp.max(
         jnp.abs(w_qarray_pad.qvalue[:, :D, :] - w_qarray_ptq.qvalue[:, :D, :])
@@ -139,7 +140,7 @@ def test_dot_general_simple(self, weight_qtype, act_qtype, tile_size):
     )
 
     w_qarray_pad = padded_ptq.quantize_act(w, how, rule, None)
-    w_qarray_ptq = ptq.quantize_act(w_padded, how, rule, None)
+    w_qarray_ptq = boxed_param.quantize_act(w_padded, how, rule, None)
 
     q_diff = jnp.max(
         jnp.abs(w_qarray_pad.qvalue[:, :D, :] - w_qarray_ptq.qvalue[:, :D, :])
@@ -211,7 +212,7 @@ def test_einsum_simple_2d_scales(self):
     )
 
     w_qarray_pad = padded_ptq.quantize_act(w, how, rule, None)
-    w_qarray_ptq = ptq.quantize_act(w_padded, how, rule, None)
+    w_qarray_ptq = boxed_param.quantize_act(w_padded, how, rule, None)
 
     q_diff = jnp.max(
         jnp.abs(w_qarray_pad.qvalue - w_qarray_ptq.qvalue[:E, :D, :F])