Add reverse-mode AD support and dimension permutation to qwix hijax.

This change implements Vector-Jacobian Product (VJP) for the hijax conversion primitives. It also introduces a PermuteDims primitive, along with permute_dims and transpose functions for HiQArray, including their VJP implementations. Finally, it adds unit tests to verify the backward passes and permutation operations.

PiperOrigin-RevId: 958434873

Author: chapman20j

qwix/contrib/hijax/convert.py

@@ -17,8 +17,12 @@
 
 # pyrefly: ignore-errors
 
+import functools
+
 import jax
 import jax.experimental.hijax as hjx
+import jax.numpy as jnp
+import numpy as np
 import qwix.contrib.hijax.hiqarray as hq
 
 
@@ -74,7 +78,7 @@ def __init__(
     self.out_aval = hq.HiQArrayTy(out_qvalue_ty, scale_ty, zero_point_ty)
     self.params = dict(
         quantize_fn=quantize_fn,
-        **quantize_kwargs,
+        quantize_kwargs=quantize_kwargs,
     )
     # For type checking
     self.quantize_fn = quantize_fn
@@ -88,6 +92,16 @@ def expand(self, data, scale, zero_point):
     )
     return hq.HiQArray(quantized_data, scale, zero_point)
 
+  # Reverse mode ad
+  def vjp_fwd(
+      self, nzs_in, data: jax.Array, scale: jax.Array, zero_point: jax.Array
+  ):
+    return self(data, scale, zero_point), None
+
+  def vjp_bwd_retval(self, res, g, /):
+    # Use Straight-Through Estimate (STE)
+    return (g, None, None)
+
 
 def to_hiqarray(
     data: jax.Array,
@@ -128,6 +142,13 @@ def expand(self, qarray: hq.HiQArray):
     )
     return dequantized_data
 
+  # Reverse mode ad
+  def vjp_fwd(self, nzs_in, qarray: hq.HiQArray):
+    return self(qarray), None
+
+  def vjp_bwd_retval(self, res, g, /):
+    return (g,)
+
 
 def from_hiqarray(
     qarray: hq.HiQArray, *, dequantize_fn, **dequantize_kwargs
@@ -138,3 +159,69 @@ def from_hiqarray(
       ty, dequantize_fn=dequantize_fn, **dequantize_kwargs
   )
   return from_qarray_instance(qarray)
+
+
+class PermuteDims(hjx.VJPHiPrimitive):
+  """Hijax primitive for permuting dimensions of a HiQArray."""
+
+  def __init__(
+      self,
+      in_aval: hq.HiQArrayTy,
+      axes: tuple[int, ...],
+  ):
+    self.in_avals = (in_aval,)
+    self.out_aval = self._permute_dims_aval(in_aval, axes)
+    self.params = dict(axes=axes)
+    # For pytype warnings
+    self.axes = axes
+    super().__init__()
+
+  # Private functions
+  @staticmethod
+  def _permute_dims_aval(
+      in_aval: hq.HiQArrayTy, axes: tuple[int, ...]
+  ) -> hq.HiQArrayTy:
+    inner_fn = functools.partial(
+        jax.eval_shape, lambda x: jnp.permute_dims(x, axes=axes)
+    )
+
+    def fn(x):
+      sds = inner_fn(x)
+      return jax._src.core._sds_aval_mapping(sds)  # pylint: disable=protected-access
+
+    lo_avals_permuted = jax.tree_util.tree_map(fn, in_aval.lo_ty())
+    return hq.HiQArrayTy.raise_ty(lo_avals_permuted)
+
+  def expand(self, qarray: hq.HiQArray):
+    return hq.HiQArray(
+        jnp.permute_dims(qarray.qvalue, self.axes),
+        jnp.permute_dims(qarray.scale, self.axes),
+        (
+            jnp.permute_dims(qarray.zero_point, self.axes)
+            if qarray.zero_point is not None
+            else None
+        ),
+    )
+
+  # Reverse mode ad
+  def vjp_fwd(self, nzs_in, qarray: hq.HiQArray):
+    return permute_dims(qarray, self.axes), None
+
+  def vjp_bwd_retval(self, res, g, /):
+    inv_perm = tuple(np.argsort(self.axes))
+    return (jnp.permute_dims(g, inv_perm),)
+
+
+def permute_dims(qarray: hq.HiQArray, axes: tuple[int, ...]) -> hq.HiQArray:
+  ty = jax.typeof(qarray)
+  permute_axes_instance = PermuteDims(ty, axes)
+  return permute_axes_instance(qarray)
+
+
+def transpose(qarray: hq.HiQArray) -> hq.HiQArray:
+  if qarray.ndim < 2:
+    raise ValueError(f"Called transpose on HiQArray of shape {qarray.shape}")
+
+  s = list(range(qarray.ndim))
+  new_s = s[:-2] + [s[-1], s[-2]]
+  return permute_dims(qarray, new_s)

tests/contrib/hijax/hiqarray_test.py

@@ -1,3 +1,5 @@
+import functools
+
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
@@ -10,11 +12,10 @@
 class HiqarrayTest(parameterized.TestCase):
 
   def test_as_hiqarray(self):
-    x = jnp.ones((16, 32, 64))
-    scale = jnp.ones((1, 1, 64))
-    zero_point = jnp.zeros((1, 1, 64))
+    x = jnp.ones((16, 32))
+    scale = jnp.ones((16, 1))
 
-    xq = convert.as_hiqarray(x, scale, zero_point)
+    xq = convert.as_hiqarray(x, scale, None)
 
     self.assertIsInstance(xq, hiqarray.HiQArray)
 
@@ -51,6 +52,77 @@ def dequantize_fn(data, scale, zero_point):
     self.assertEqual(y.shape, x.shape)
     self.assertEqual(y.dtype, x.dtype)
 
+  def test_to_hiqarray_bwd(self):
+    key = jax.random.key(0)
+    key1, key2 = jax.random.split(key, 2)
+    x = jax.random.normal(key1, (16, 32))
+    scale = jax.random.normal(key2, (16, 1))
+
+    def quantize_fn(data, scale, zero_point, qtype):
+      return qarray.quantize_with_scale_zero_point(
+          data, qtype, scale, zero_point
+      ).qvalue
+
+    fn = functools.partial(
+        convert.to_hiqarray, quantize_fn=quantize_fn, qtype=jnp.int8
+    )
+
+    _, vjp_fn = jax.vjp(fn, x, scale, None)
+    g = jnp.ones_like(x)
+    bwd = vjp_fn(g)
+
+    self.assertTrue(jnp.allclose(bwd[0], g))
+    self.assertLess(jnp.max(jnp.abs(bwd[1])), 1e-6)
+    self.assertIsNone(bwd[2])
+
+  def test_from_hiqarray_bwd(self):
+    x = jnp.ones((16, 32))
+    scale = jnp.ones((16, 1))
+
+    xq = convert.as_hiqarray(x, scale, None)
+
+    def dequantize_fn(data, scale, zero_point):
+      del zero_point  # unused
+      return qarray.call_with_generic_broadcast(jnp.multiply, data, scale)
+
+    fn = functools.partial(convert.from_hiqarray, dequantize_fn=dequantize_fn)
+
+    _, vjp_fn = jax.vjp(fn, xq)
+    g = jnp.ones_like(xq.qvalue)
+    bwd = vjp_fn(g)
+
+    self.assertTrue(jnp.allclose(bwd[0], g))
+
+  def test_permute_dims(self):
+    key = jax.random.key(0)
+    key1, key2 = jax.random.split(key, 2)
+    x = jax.random.normal(key1, (16, 32))
+    scale = jax.random.normal(key2, (16, 1))
+
+    xq = convert.as_hiqarray(x, scale, None)
+
+    yq = convert.permute_dims(xq, (1, 0))
+    zq = convert.transpose(xq)
+
+    self.assertTrue(jnp.allclose(jnp.transpose(x), yq.qvalue))
+    self.assertTrue(jnp.allclose(jnp.transpose(scale), yq.scale))
+    self.assertTrue(jnp.allclose(yq.qvalue, zq.qvalue))
+    self.assertTrue(jnp.allclose(yq.scale, zq.scale))
+
+  def test_permute_dims_bwd(self):
+    key = jax.random.key(0)
+    key1, key2, key3 = jax.random.split(key, 3)
+    x = jax.random.normal(key1, (16, 32))
+    scale = jax.random.normal(key2, (16, 1))
+
+    xq = convert.as_hiqarray(x, scale, None)
+
+    _, vjp_fn = jax.vjp(convert.transpose, xq)
+    g = jax.random.normal(key3, (32, 16))
+    bwd = vjp_fn(g)
+
+    self.assertTrue(jnp.allclose(bwd[0], jnp.transpose(g)))
+
 
 if __name__ == "__main__":
   absltest.main()