Make HowToQuantize fields optional with default values.

This simplifies the usage of `qarray.HowToQuantize` by providing default values for `channelwise_axes`, `tiled_axes`, and `calibration_method`. Tests are updated to reflect this change.

PiperOrigin-RevId: 811189153

Author: liudangyi

qwix/_src/core/qarray.py

@@ -241,17 +241,19 @@ class HowToQuantize:
   qtype: jax.typing.DTypeLike
   # Channelwise axes will have individual scales, which has the same effect
   # as setting their tile sizes to 1 in tiled_axes.
-  channelwise_axes: Collection[int]
+  channelwise_axes: Collection[int] = ()
   # Tiled axes have subchannel quantization enabled. The value is a mapping
   # from the tiled axis to the tile size. If the tile size is a float, it has
   # to be "1 / tile_count" and the actual tile size will be
   # round(axis_size * tile_size). Note that 1 and 1.0 have very different
   # meanings: a tile size of 1 means to use per-channel scale, while a
   # tile size of 1.0 means to use shared scale.
-  tiled_axes: Mapping[int, int | float]
+  tiled_axes: Mapping[int, int | float] = dataclasses.field(
+      default_factory=dict
+  )
   # The calibration method to use. The format is <method>[,<args>], e.g.
   # "absmax" or "fixed,-10,10". Check calibrate() for supported methods.
-  calibration_method: str
+  calibration_method: str = 'absmax'
 
 
 ShapeT: TypeAlias = Sequence[int]

tests/core/einsum_test.py

@@ -275,12 +275,7 @@ def test_fake_quantization(self):
     fp_res = jnp.einsum(
         einsum_str, lhs, rhs, precision=jax.lax.Precision.HIGHEST
     )
-    how = qarray.HowToQuantize(
-        qtype=jnp.int8,
-        channelwise_axes=(),
-        tiled_axes={},
-        calibration_method='absmax',
-    )
+    how = qarray.HowToQuantize(qtype=jnp.int8)
     lhs = qarray.quantize(lhs, how)
     rhs = qarray.quantize(rhs, how)
 
@@ -321,21 +316,11 @@ def test_dequant_on_inputs(self):
     rhs = self._make_array((128, 128, 16), jnp.bfloat16)
     lhs = qarray.quantize(
         lhs,
-        qarray.HowToQuantize(
-            qtype=jnp.int8,
-            channelwise_axes=(0, 1),
-            tiled_axes={},
-            calibration_method='absmax',
-        ),
+        qarray.HowToQuantize(qtype=jnp.int8, channelwise_axes=(0, 1)),
     )
     rhs = qarray.quantize(
         rhs,
-        qarray.HowToQuantize(
-            qtype=jnp.int8,
-            channelwise_axes=(0, 2),
-            tiled_axes={},
-            calibration_method='absmax',
-        ),
+        qarray.HowToQuantize(qtype=jnp.int8, channelwise_axes=(0, 2)),
     )
     out = einsum.einsum('TNH,NHD -> TD', lhs, rhs)
     self.assertEqual(out.shape, (16, 16))

tests/core/pallas_test.py

@@ -211,9 +211,7 @@ def dequantize_pallas(q: qarray.QArray):
     x = jax.random.uniform(jax.random.key(0), input_shape, jnp.float32)
     how = qarray.HowToQuantize(
         qtype="int8",
-        channelwise_axes=[],
         tiled_axes=tiled_axes,
-        calibration_method="absmax",
     )
     qx = qarray.quantize(x, how)
     self.assertTrue(jnp.allclose(dequantize_pallas(qx), qarray.dequantize(qx)))
@@ -262,7 +260,6 @@ def pallas_batch_matmul(
         qtype="int8",
         channelwise_axes=[0, 1],
         tiled_axes={2: 128},
-        calibration_method="absmax",
     )
     qx = qarray.quantize(
         jax.random.uniform(jax.random.key(0), (4, 256, 256), jnp.float32), x_how
@@ -271,7 +268,6 @@ def pallas_batch_matmul(
         qtype="int8",
         channelwise_axes=[2],
         tiled_axes={1: 128},
-        calibration_method="absmax",
     )
     qy = qarray.quantize(
         jax.random.uniform(jax.random.key(1), (4, 256, 256), jnp.float32), y_how

tests/core/qarray_test.py

@@ -153,12 +153,7 @@ def test_exact_quantization(self, with_error):
       array += jax.random.uniform(
           jax.random.key(42), array.shape, minval=-1e-7, maxval=1e-7
       )
-    how = qarray.HowToQuantize(
-        qtype=jnp.int8,
-        channelwise_axes=[],
-        tiled_axes={},
-        calibration_method='minmax',
-    )
+    how = qarray.HowToQuantize(qtype=jnp.int8, calibration_method='minmax')
     q_array = qarray.quantize(array, how)
     self.assertEqual(q_array.zero_point, jnp.array(-128, dtype=jnp.int8), array)
     expected_q_array = jnp.arange(-128, 128, dtype=jnp.int8)