TL;DR

• Subclass HiType and implement lo_ty, lower_val, and raise_val to say how the type and its values lower to ordinary (“lojax”) arrays, then call register_hitype to associate your value class with your type.

• Write VJPHiPrimitive subclasses whose in_avals/out_aval mention the new type; these are the only way values of the type get produced and consumed.

• For autodiff, implement to_tangent_aval on the type, and VJP/JVP rules on the primitives.

• For vmap, implement dec_rank and inc_rank on the type along with a MappingSpec subclass of your own design, and batch rules on the primitives. Mapped-over hi type arguments require an explicit axis_size and spec-valued in_axes/out_axes entries.

• For sharding in types (explicit mode), record sharding data on your type (e.g. a NamedSharding field), consume it in lo_ty, and propagate it in your primitives’ typing rules.

Example: quantized arrays

Say we want to work with arrays quantized to int8. A quantized array is really a pair of arrays: the int8 values, and a floating point scale shared by each row (that is, we quantize along the last axis, one scale per row, as in common per-row/per-channel quantization schemes):

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
# (8 CPU devices, for the sharding sections at the end)

from dataclasses import dataclass

import jax
import jax.numpy as jnp

@dataclass(frozen=True)
class QArray:
  qvalue: jax.Array   # int8[*leading, n]
  scale: jax.Array    # f32[*leading]

We could register QArray as a pytree and be done. But consider what we’d give up:

• Invariants. The two components are coupled: scale must have the shape of qvalue minus its last axis, and qvalue is only meaningful together with its scale. As a pytree, nothing stops code from crossing the streams; under transformations, JAX itself sees only independent leaves.

• Types in jaxprs. As a pytree, a quantized array appears in traced code as two unrelated array values. We’d rather see one value, of one type, so jaxprs say what they mean.

• Tangents. A quantized array’s values live on a discrete grid, so it makes no sense to perturb them along the grid. But a pytree’s tangent type is forced to be the pytree of its leaves’ tangent types — and the tangent type of an integer array like qvalue is a float0 array, which can only carry a trivial payload. So as a pytree, a quantized array would admit no useful perturbations at all. What we want is to choose a tangent type for the quantized array as a whole, such as the continuous f32 arrays that the quantized values approximate.

So instead we’ll make QArray a hijax type.

The type

A hijax type is a subclass of HiType. The required core is small:

• lo_ty says which lojax (array) types make up the type;

• lower_val and raise_val convert values to and from that list of arrays;

• the type must be hashable and comparable for equality (a frozen dataclass gives us both).

This is like the pytree flatten/unflatten interface, but it lives at the level of types: given only the type, JAX can compute the lowered types, without needing a value in hand.

We also give the type a sharding field, recording how values are partitioned across devices; it can be ignored until the sharding sections near the end of this document. We reuse JAX’s NamedSharding to describe the partitioning of the qvalue component, and derive scale’s partitioning from it by dropping the last axis. Nothing about the field is special to JAX, which never interprets it: only our own methods consume it, chiefly lo_ty, which stamps the component types with their shardings. You’re free to track sharding information on your type with an object of your own design instead, so long as you consume it the same way.

from jax.experimental.hijax import HiType, ShapedArray, register_hitype
from jax.sharding import NamedSharding

@dataclass(frozen=True)
class QArrayTy(HiType):
  shape: tuple[int, ...]
  sharding: NamedSharding  # qvalue's sharding; scale's is derived from it

  # lowering: which array types make up this type, and how values convert
  def lo_ty(self):
    scale_sharding = self.sharding.update(spec=jax.P(*self.sharding.spec[:-1]))
    return [ShapedArray(self.shape, jnp.dtype('int8'),
                        sharding=self.sharding),
            ShapedArray(self.shape[:-1], jnp.dtype('float32'),
                        sharding=scale_sharding)]
  def lower_val(self, q):
    return [q.qvalue, q.scale]
  def raise_val(self, qvalue, scale):
    return QArray(qvalue, scale)

  # autodiff: tangents of quantized arrays are plain float arrays (see below)
  def to_tangent_aval(self):
    return ShapedArray(self.shape, jnp.dtype('float32'),
                       sharding=self.sharding)

  # printing, e.g. in jaxprs
  def str_short(self, short_dtypes=False, mesh_axis_types=False):
    dims = [str(d) if p is None else f'{d}@{p}'
            for d, p in zip(self.shape, self.sharding.spec)]
    return f'q8[{",".join(dims)}]'
  __repr__ = str_short

register_hitype(QArray, lambda q: QArrayTy(q.qvalue.shape,
                                           jax.typeof(q.qvalue).sharding))

The register_hitype call associates the value class with the type: its second argument computes the type of any given value, analogous to how jax.typeof maps an array to its ShapedArray type. (Ours reads both the shape and the sharding off the qvalue component — every array carries a sharding, trivial when no mesh is in play.) Indeed after registration, jax.typeof works on QArrays, and JAX transformations accept them anywhere a value is expected.

The primitives

With a pytree, users construct and take apart values freely. With a hijax type, values are produced and consumed only by hijax primitives whose declared types mention the new type. That’s where invariants get enforced: if every primitive preserves them, they always hold.

Our two primitives are quantize and dequantize, written with the VJPHiPrimitive API from Custom derivative rules with hijax primitives. Each declares its input and output types, gives its implementation in expand, and (looking ahead to autodiff) carries a straight-through-estimator VJP rule:

from jax.experimental.hijax import VJPHiPrimitive

class Quantize(VJPHiPrimitive):
  def __init__(self, x_aval):
    if x_aval.dtype != jnp.dtype('float32'): raise TypeError(x_aval.dtype)
    self.in_avals = (x_aval,)
    self.out_aval = QArrayTy(x_aval.shape, x_aval.sharding)
    self.params = {}
    super().__init__()

  def expand(self, x):
    scale = jnp.max(jnp.abs(x), axis=-1) / 127.
    qvalue = jnp.round(x / scale[..., None]).astype(jnp.int8)
    return QArray(qvalue, scale)

  # straight-through estimator: differentiate as if it's the identity
  def vjp_fwd(self, nzs_in, x):
    return self(x), None

  def vjp_bwd_retval(self, _res, g):
    return (g,)

class Dequantize(VJPHiPrimitive):
  def __init__(self, q_aval):
    self.in_avals = (q_aval,)
    self.out_aval = ShapedArray(q_aval.shape, jnp.dtype('float32'),
                                sharding=q_aval.sharding)
    self.params = {}
    super().__init__()

  def expand(self, qx):
    return qx.qvalue.astype('float32') * qx.scale[..., None]

  def vjp_fwd(self, nzs_in, qx):
    return self(qx), None

  def vjp_bwd_retval(self, _res, g):
    return (g,)

def quantize(x):
  return Quantize(jax.typeof(x))(x)

def dequantize(qx):
  return Dequantize(jax.typeof(qx))(qx)

Notice that Quantize’s out_aval and Dequantize’s in_avals are QArrayTys: the new type appears in primitive type signatures just like array types do. Also notice expand freely constructs and inspects the QArray value class; primitive implementations are inside the abstraction boundary.

Everything works eagerly:

x = jnp.array([[1., 2., 3.],
               [4., -5., 6.]])

qx = quantize(x)
print(qx)
print(jax.typeof(qx))
print(dequantize(qx))

QArray(qvalue=Array([[  42,   85,  127],
       [  85, -106,  127]], dtype=int8), scale=Array([0.02362205, 0.04724409], dtype=float32))
q8[2,3]
[[ 0.992126  2.007874  3.      ]
 [ 4.015748 -5.007874  6.      ]]

try:
  jax.jit(lambda qx: qx.qvalue)(qx)
except AttributeError as e:
  print('AttributeError:', e)

AttributeError: DynamicJaxprTracer has no attribute qvalue

def bad_quantize(x):
  scale = jnp.max(jnp.abs(x), axis=-1) / 127.
  return QArray(jnp.round(x / scale[..., None]).astype('int8'), scale)

try:
  jax.jit(bad_quantize)(x)
except TypeError as e:
  print('TypeError:', e)

TypeError: No constant handler for type: <class 'jax._src.interpreters.partial_eval.DynamicJaxprTracer'>

class MatmulQ(VJPHiPrimitive):
  def __init__(self, x_aval, q_aval):
    if not (isinstance(q_aval, QArrayTy) and len(x_aval.shape) == 2 and
            len(q_aval.shape) == 2 and x_aval.shape[1] == q_aval.shape[0]):
      raise TypeError(f'bad matmul_q operand types: {x_aval} @ {q_aval}')
    self.in_avals = (x_aval, q_aval)
    x_spec, q_spec = x_aval.sharding.spec, q_aval.sharding.spec
    if x_spec[1] is not None or q_spec[0] is not None:
      raise TypeError('matmul_q requires unsharded contraction axes, got '
                      f'{x_aval} @ {q_aval}')
    out_sharding = x_aval.sharding.update(spec=jax.P(x_spec[0], q_spec[1]))
    self.out_aval = ShapedArray((x_aval.shape[0], q_aval.shape[1]),
                                jnp.dtype('float32'), sharding=out_sharding)
    self.params = {}
    super().__init__()

  def expand(self, x, qw):
    # fold the per-row scales into the dense operand, then apply one matmul
    # directly against the int8 payload
    return (x * qw.scale) @ qw.qvalue.astype(jnp.float32)

  def vjp_fwd(self, nzs_in, x, qw):
    return self(x, qw), (x, qw)

  def vjp_bwd_retval(self, res, g):
    x, qw = res
    w = dequantize(qw)   # rules are traced code: use primitives here
    # cotangents live where the primals live, so use the primal operands'
    # shardings to disambiguate the (possibly sharded) contractions
    return (jnp.matmul(g, w.T, out_sharding=jax.typeof(x).sharding),
            jnp.matmul(x.T, g, out_sharding=jax.typeof(w).sharding))

def matmul_q(x, qw):
  return MatmulQ(jax.typeof(x), jax.typeof(qw))(x, qw)

w = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.
qw = quantize(w)

print(matmul_q(x, qw))
print(x @ dequantize(qw))  # reference

[[2.6627297 3.1706038 3.6587927 4.166667 ]
 [2.3077428 2.7447505 3.1463253 3.5833333]]
[[2.6627297 3.1706038 3.6587927 4.166667 ]
 [2.3077426 2.7447505 3.1463253 3.5833333]]

Hi types in jaxprs

When we trace, the quantized array appears as a single value of type q8[2,3], produced by one equation and consumed by another:

jax.jit(lambda x: dequantize(quantize(x))).trace(x).jaxpr

{ lambda ; a:f32[2,3]. let
    b:q8[2,3] = call_hi_primitive[_prim=Quantize[{}]] a
    c:f32[2,3] = call_hi_primitive[_prim=Dequantize[{}]] b
  in (c,) }

Compare to the pytree approach, where the same computation would show four array-typed intermediates with no indication that they pair up. The hi type only disappears at lowering time, when expand is traced and each q8[...]-typed value is expanded into the array components given by lo_ty.

Ops with mixed operand kinds read just as directly — one equation with a f32[2,3] @ q8[3,4] -> f32[2,4] signature:

jax.jit(matmul_q).trace(x, qw).jaxpr
#
# `jit` works, with quantized arrays as arguments, results, and
# intermediates:

{ lambda ; a:f32[2,3] b:q8[3,4]. let
    c:f32[2,4] = call_hi_primitive[_prim=MatmulQ[{}]] a b
  in (c,) }

print(jax.jit(lambda x: dequantize(quantize(x)))(x))   # QArray internal

qx2 = jax.jit(quantize)(x)                             # QArray result
print(jax.typeof(qx2))

print(jax.jit(dequantize)(qx2))                        # QArray argument

[[ 0.992126  2.007874  3.      ]
 [ 4.015748 -5.007874  6.      ]]
q8[2,3]
[[ 0.992126  2.007874  3.      ]
 [ 4.015748 -5.007874  6.      ]]

Autodiff and tangent types

Here’s where hi types earn their keep. On the type, we implemented

  def to_tangent_aval(self):
    return ShapedArray(self.shape, jnp.dtype('float32'))

which says: the tangent type of a quantized array is a plain f32 array. No pytree can express this: a pytree’s tangent type is always the pytree of its leaves’ tangent types, and for the int8 leaf qvalue that means a trivial float0 tangent.

Together with the straight-through VJP rules on the primitives, gradients flow through quantization as if it were the identity:

def f(x):
  return jnp.sum(dequantize(quantize(x)))

And differentiating with respect to a quantized array input produces a plain float array, as the tangent type dictates:

def g(qx):
  return jnp.sum(dequantize(qx) ** 2)

print(jax.grad(g)(qx))
print(jax.typeof(jax.grad(g)(qx)))

[[  1.984252   4.015748   6.      ]
 [  8.031496 -10.015748  12.      ]]
float32[2,3]

The same goes for ops with mixed operand kinds, like the quantized matmul: differentiating a loss with respect to both operands gives an f32 gradient for the dense activations and an f32 gradient for the quantized weights:

def loss(x, qw):
  return jnp.sum(matmul_q(x, qw) ** 2)

grad_x, grad_qw = jax.grad(loss, argnums=(0, 1))(x, qw)
print(jax.typeof(grad_x), jax.typeof(grad_qw))

float32[2,3] float32[3,4]

Notice that making the tangent type an f32 array was a choice, and there’s a real design space here. We could instead have made the tangent type of QArrayTy be QArrayTy itself, so that tangents and cotangents are quantized too — a different tradeoff, sensible for different applications. (For that choice, since the tangent type is then a hi type, we’d also implement vspace_zero and vspace_add on it so autodiff can instantiate and accumulate cotangents.) This flexibility is why hi types are a user extension point: for each piece of JAX — tracing, lowering, autodiff, and batching — you set up how your type participates, however your situation needs.

vmap and mapping specs

What does it mean to map over a quantized array? For arrays, vmap’s in_axes and out_axes are axis indices, and JAX can infer the mapped axis size from the argument’s shape. For a general hi type, JAX doesn’t guess: you define a “mapping spec” type that says how values of your type are mapped, users pass instances of it as in_axes/out_axes entries, and they pass axis_size explicitly when it can’t be inferred from an array argument.

For our quantized arrays, thanks to the per-row scales, a batch of QArrays is just a bigger QArray: stacking n quantized arrays of type q8[2,3] along a new leading axis gives a q8[n,2,3], with qvalue of shape (n, 2, 3) and scale of shape (n, 2). So the only mapping notion we need is “the leading axis,” and our spec type doesn’t need to carry any data at all:

from jax.experimental.hijax import MappingSpec

@dataclass(frozen=True)
class QArraySpec(MappingSpec):
  pass  # QArrays are only mapped along their leading axis

(Specs can be as rich as your type demands. A tuple-like hi type might use a spec carrying one axis per component; see the tuple example at the end of this document.)

On the type, we implement dec_rank and inc_rank, the hi type analogues of “remove the mapped axis” and “add the mapped axis.” They take the axis size and a spec, and return the element type and the batched type, respectively:

def qarray_dec_rank(self, size, spec):
  assert isinstance(spec, QArraySpec) and self.shape[0] == size
  return QArrayTy(self.shape[1:],
                  self.sharding.update(spec=jax.P(*self.sharding.spec[1:])))

def qarray_inc_rank(self, size, spec):
  assert isinstance(spec, QArraySpec)
  return QArrayTy((size, *self.shape),
                  self.sharding.update(spec=jax.P(None, *self.sharding.spec)))

QArrayTy.dec_rank = qarray_dec_rank
QArrayTy.inc_rank = qarray_inc_rank

(We’re attaching methods to the class as we go, notebook-style; in real code these would just be more methods in the class QArrayTy definition.)

On the primitives, we implement batch rules. A batch rule receives the batched arguments along with their mapping specs (None for unbatched arguments, an integer axis for batched array arguments, and a spec instance for batched hi type arguments), and returns the batched result along with its mapping spec. Note that a rule should be prepared for any combination of batched and unbatched arguments:

def quantize_batch(self, axis_data, args, in_dims):
  x, = args
  d, = in_dims
  if d is None:
    return quantize(x), None
  x = jnp.moveaxis(x, d, 0)
  return quantize(x), QArraySpec()
Quantize.batch = quantize_batch

def dequantize_batch(self, axis_data, args, in_dims):
  qx, = args
  d, = in_dims
  if d is None:
    return dequantize(qx), None
  assert isinstance(d, QArraySpec)
  return dequantize(qx), 0
Dequantize.batch = dequantize_batch

Because per-row quantization applies at any rank, both rules can just apply the unbatched operation to the stacked value — the hallmark of a type whose batches are values of the same type family.

Now we can vmap. Mapping to a quantized array output, the axis size is inferred from the array argument as usual, and we pass a spec for out_axes:

xs = jnp.arange(24., dtype='float32').reshape(4, 2, 3)

qxs = jax.vmap(quantize, out_axes=QArraySpec())(xs)
print(jax.typeof(qxs))
print(qxs.qvalue.shape, qxs.scale.shape)

q8[4,2,3]
(4, 2, 3) (4, 2)

Mapping over a quantized array input, we pass a spec for in_axes — and since there’s no array argument to infer the axis size from, we must pass axis_size explicitly:

xs_roundtrip = jax.vmap(dequantize, in_axes=QArraySpec(), axis_size=4)(qxs)
print(jax.typeof(xs_roundtrip))

All the usual compositions work — vmap of jit,

print(jax.typeof(jax.vmap(jax.jit(dequantize), in_axes=QArraySpec(),
                          axis_size=4)(qxs)))

vmap of grad, and so on:

def norm_quantized(x):
  return jnp.sum(dequantize(quantize(x)) ** 2)

print(jax.vmap(jax.grad(norm_quantized))(xs).shape)

scan and the leading axis

jax.lax.scan can loop over a stacked hi value, consuming one slice per step — and it can carry and produce hi values too. Where vmap asks the user for a mapping spec, scan always walks the leading axis, so it instead asks the type: the one extra method to implement is leading_axis_spec, which returns the mapping spec describing your type’s leading axis. The dec_rank and inc_rank methods from the vmap section do the rest of the work.

def qarray_leading_axis_spec(self):
  return QArraySpec()

QArrayTy.leading_axis_spec = qarray_leading_axis_spec

Scanning over the stack of quantized arrays from the previous section, the body sees one q8[2,3] per step. As with vmap’s axis_size, when all the scanned-over values are hi types there’s no leading-axis size to infer, so we pass length explicitly (if any scanned-over value is an array, scan infers the length from it and length can be omitted):

def sum_dequantized(total, qx):
  return total + jnp.sum(dequantize(qx)), ()

total, () = jax.lax.scan(sum_dequantized, 0., qxs, length=4)
print(total)

Hi values also work as stacked outputs and as the loop carry. Here the carry is a quantized array, re-quantized after each accumulation step, and the second output stacks one fresh QArray per step into a q8[4,2,3]:

def accum_quantized(qtotal, x):
  return quantize(dequantize(qtotal) + x), quantize(2 * x)

qzero = quantize(jnp.zeros((2, 3), 'float32'))
qtotal, qys = jax.lax.scan(accum_quantized, qzero, xs)
print(jax.typeof(qtotal))
print(jax.typeof(qys))

Sharding in types: explicit mode

Finally, sharding. In JAX’s explicit sharding mode (see the parallelism guide), shardings are part of array types: jax.typeof reports how a value is partitioned across the mesh, sharding propagation happens while tracing, and mismatches surface as type errors. The sharding field on QArrayTy, and the typing rules on our primitives that propagate it, are exactly what let hi types participate. What does it mean to partition a quantized array across devices? The components move together: if we shard the rows, each device holds its rows’ quantized values and their scales. As usual there’s a design choice to make, and we made it back in lo_ty: qvalue carries the type’s sharding, and scale shards like it with the last axis dropped, so every row travels with its scale.

Let’s see it work. We make a mesh (this is where we use the 8 CPU devices requested in this document’s first cell), shard some rows across it, and quantize — the shardings propagate through our typing rules into the result type, which jax.typeof displays with @ markers:

mesh = jax.make_mesh((4,), ('i',))
jax.set_mesh(mesh)

rows = jax.device_put(jnp.arange(24., dtype='float32').reshape(8, 3),
                      jax.P('i'))

qrows = quantize(rows)
print(jax.typeof(qrows))
print(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)
print(jax.typeof(dequantize(qrows)))

q8[8@i,3]
P('i', None) P('i',)
float32[8@i,3]

The same holds while tracing — hi types in jaxprs now display their shardings, computed by our out_aval rules:

jax.jit(lambda x: dequantize(quantize(x))).trace(rows).jaxpr

{ lambda ; a:f32[8@i,3]. let
    b:q8[8@i,3] = call_hi_primitive[_prim=Quantize[{}]] a
    c:f32[8@i,3] = call_hi_primitive[_prim=Dequantize[{}]] b
  in (c,) }

The quantized matmul propagates shardings too: its typing rule partitions output rows like x’s rows and output columns like qw’s columns (requiring the contracted axes to be unsharded, since neither operand can say how a sharded contraction should land):

w2 = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.

print(jax.typeof(matmul_q(rows, quantize(w2))))                # rows sharded

qw2 = quantize(jax.device_put(w2, jax.P(None, 'i')))
print(jax.typeof(qw2))                                         # cols sharded
print(jax.typeof(matmul_q(jnp.ones((2, 3), 'float32'), qw2)))

float32[8@i,4]
q8[3,4@i]
float32[2,4@i]

And the contraction check fires as a type error, at trace time:

xk = jax.device_put(jnp.ones((2, 8), 'float32'), jax.P(None, 'i'))
qk = quantize(jax.device_put(jnp.ones((8, 3), 'float32'), jax.P('i')))

try:
  matmul_q(xk, qk)
except TypeError as e:
  print('TypeError:', e)

TypeError: matmul_q requires unsharded contraction axes, got float32[2,8@i] @ q8[8@i,3]

Autodiff composes with all of this. Recall that MatmulQ’s backward rule passed out_sharding hints to its matmuls: that’s because the cotangent for qw contracts over the row axis, which may be sharded (as it is here), and explicit mode refuses to guess how an all-sharded contraction should land. The right answer is the primal operand’s sharding — cotangents live where their primals live:

def qloss(x, qw):
  return jnp.sum(matmul_q(x, qw) ** 2)

grad_rows, grad_qw2 = jax.grad(qloss, argnums=(0, 1))(rows, quantize(w2))
print(jax.typeof(grad_rows), jax.typeof(grad_qw2))

float32[8@i,3] float32[3,4]

One caution: JAX does not cross-check a hi primitive’s declared output sharding against what its expand actually produces. The declared out_aval is what downstream code sees while tracing; expand determines what happens at runtime. Keeping them consistent is part of the typing rule’s job.

shard_map and partition specs

Explicit mode partitions values while keeping one global view of the program. Its complement is shard_map, which gives a per-device view. To cross that boundary, a quantized array needs one more piece: a partition spec type of our own, saying how shard_map’s in_specs/out_specs apply to it. We keep the same design as above: a quantized array is sharded along its leading axes only.

For shard_map, we express this with an HiPspec subclass — the partition spec analogue of the MappingSpec above. Users pass instances of it as in_specs/out_specs entries, and its to_lo method says how it translates to one jax.P partition spec per lowered component (in lo_ty order):

from jax.experimental.hijax import HiPspec

@dataclass(frozen=True)
class QArrayP(HiPspec):
  spec: jax.P  # partitioning of the leading axes; the last axis stays whole

  def to_lo(self):
    return (self.spec, self.spec)  # qvalue and scale shard together

(Since scale has one axis fewer than qvalue, handing both components the same partition spec says exactly that the leading axes shard together while the trailing axis of qvalue is untouched.)

On the type, shard and unshard compute the per-device shard type from the global type and vice versa, delegating to the component types:

def qarray_shard(self, mesh, manual_axes, check_vma, spec):
  qvalue_ty, _ = self.lo_ty()
  qspec, _ = spec.to_lo()
  shard_ty = qvalue_ty.shard(mesh, manual_axes, check_vma, qspec)
  return QArrayTy(shard_ty.shape, shard_ty.sharding)

def qarray_unshard(self, mesh, check_vma, spec):
  qvalue_ty, _ = self.lo_ty()
  qspec, _ = spec.to_lo()
  full_ty = qvalue_ty.unshard(mesh, check_vma, qspec)
  return QArrayTy(full_ty.shape, full_ty.sharding)

QArrayTy.shard = qarray_shard
QArrayTy.unshard = qarray_unshard

Now quantized arrays can cross shard_map boundaries in either direction, continuing with the mesh and rows from the previous section. Producing a sharded quantized array:

@jax.jit
@jax.shard_map(in_specs=jax.P('i'), out_specs=QArrayP(jax.P('i')))
def quantize_shards(x):
  assert jax.typeof(x).shape == (2, 3)   # each device sees two rows
  return quantize(x)

qrows = quantize_shards(rows)
print(jax.typeof(qrows))
print(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)

q8[8@i,3]
P('i', None) P('i',)

And consuming one, where each device sees a per-shard QArray of its own rows:

@jax.jit
@jax.shard_map(in_specs=QArrayP(jax.P('i')), out_specs=jax.P('i'))
def dequantize_shards(qx):
  assert jax.typeof(qx).shape == (2, 3)  # a per-device QArray shard
  return dequantize(qx)

print(jnp.max(jnp.abs(dequantize_shards(qrows) - rows)))

Because scales are per-row, quantizing shard-by-shard agrees exactly with quantizing globally — the same property that made batching pleasant:

qrows_global = quantize(rows)
assert (qrows.qvalue == qrows_global.qvalue).all()
assert (qrows.scale == qrows_global.scale).all()

(For autodiff through a shard_map, there’s a bit more to implement: to_tangent_spec and to_ct_spec on the spec type, and nospec on the hi type, which is used to shard autodiff residuals.)

Additional examples

The recipe is always the same — a value class, a HiType with the lo_ty/lower_val/raise_val lowering triple, and primitives whose type signatures mention the new type — so further examples can be read quickly.

@dataclass(frozen=True)
class Rank1:
  col: jax.Array   # f32[m]
  row: jax.Array   # f32[n]

@dataclass(frozen=True)
class Rank1Ty(HiType):
  shape: tuple[int, int]   # (m, n), the dense shape represented
  sharding: NamedSharding  # sharding of the dense shape; the factors' derive

  def lo_ty(self):
    (m, n), spec = self.shape, self.sharding.spec
    return [ShapedArray((m,), jnp.dtype('float32'),
                        sharding=self.sharding.update(spec=jax.P(spec[0]))),
            ShapedArray((n,), jnp.dtype('float32'),
                        sharding=self.sharding.update(spec=jax.P(spec[1])))]
  def lower_val(self, r1):
    return [r1.col, r1.row]
  def raise_val(self, col, row):
    return Rank1(col, row)

  # rank-1 matrices aren't closed under addition, so tangents are dense
  def to_tangent_aval(self):
    return ShapedArray(self.shape, jnp.dtype('float32'),
                       sharding=self.sharding)

  def str_short(self, short_dtypes=False, mesh_axis_types=False):
    dims = [str(d) if p is None else f'{d}@{p}'
            for d, p in zip(self.shape, self.sharding.spec)]
    return f'r1[{",".join(dims)}]'
  __repr__ = str_short

def typeof_rank1(r1):
  col_s, row_s = jax.typeof(r1.col).sharding, jax.typeof(r1.row).sharding
  sharding = col_s.update(spec=jax.P(col_s.spec[0], row_s.spec[0]))
  return Rank1Ty((r1.col.shape[0], r1.row.shape[0]), sharding)

register_hitype(Rank1, typeof_rank1)

class Outer(VJPHiPrimitive):
  def __init__(self, col_aval, row_aval):
    if not (len(col_aval.shape) == 1 and len(row_aval.shape) == 1):
      raise TypeError(f'bad outer factor types: {col_aval}, {row_aval}')
    sharding = col_aval.sharding.update(
        spec=jax.P(col_aval.sharding.spec[0], row_aval.sharding.spec[0]))
    self.in_avals = (col_aval, row_aval)
    self.out_aval = Rank1Ty((col_aval.shape[0], row_aval.shape[0]), sharding)
    self.params = {}
    super().__init__()

  def expand(self, col, row):
    return Rank1(col, row)

  def vjp_fwd(self, nzs_in, col, row):
    return self(col, row), (col, row)

  def vjp_bwd_retval(self, res, g):   # g is dense, f32[m, n]
    col, row = res
    return g @ row, col @ g

class Factors(VJPHiPrimitive):
  def __init__(self, r1_aval):
    self.in_avals = (r1_aval,)
    self.out_aval = tuple(r1_aval.lo_ty())  # the factor types are the lo types
    self.params = {}
    super().__init__()

  def expand(self, r1):
    return (r1.col, r1.row)

class MatmulR1(VJPHiPrimitive):
  def __init__(self, x_aval, r1_aval):
    if not (isinstance(r1_aval, Rank1Ty) and len(x_aval.shape) == 2 and
            x_aval.shape[1] == r1_aval.shape[0]):
      raise TypeError(f'bad matmul_r1 operand types: {x_aval} @ {r1_aval}')
    self.in_avals = (x_aval, r1_aval)
    out_sharding = x_aval.sharding.update(
        spec=jax.P(x_aval.sharding.spec[0], r1_aval.sharding.spec[1]))
    self.out_aval = ShapedArray((x_aval.shape[0], r1_aval.shape[1]),
                                jnp.dtype('float32'), sharding=out_sharding)
    self.params = {}
    super().__init__()

  def expand(self, x, r1):
    return jnp.outer(x @ r1.col, r1.row)   # never materialize col ⊗ row

  def vjp_fwd(self, nzs_in, x, r1):
    col, row = factors(r1)  # rules are traced code: primitives, not attributes
    return self(x, r1), (x, col, row)

  def vjp_bwd_retval(self, res, g):
    x, col, row = res
    return jnp.outer(g @ row, col), x.T @ g

def outer(col, row):
  return Outer(jax.typeof(col), jax.typeof(row))(col, row)

def factors(r1):
  return Factors(jax.typeof(r1))(r1)

def matmul_r1(x, r1):
  return MatmulR1(jax.typeof(x), jax.typeof(r1))(x, r1)

col = jnp.arange(6., dtype='float32') / 6.
row = jnp.arange(5., dtype='float32') / 5.
r1 = outer(col, row)
print(jax.typeof(r1))

acts = jnp.ones((3, 6), 'float32')
print(jnp.max(jnp.abs(matmul_r1(acts, r1) - acts @ jnp.outer(col, row))))

jax.jit(matmul_r1).trace(acts, r1).jaxpr

{ lambda ; a:f32[3,6] b:r1[6,5]. let
    c:f32[3,5] = call_hi_primitive[_prim=MatmulR1[{}]] a b
  in (c,) }

def r1_loss(x, r1):
  return jnp.sum(matmul_r1(x, r1) ** 2)

print(jax.typeof(jax.grad(r1_loss, argnums=1)(acts, r1)))

def factor_loss(col, row):
  return jnp.sum(matmul_r1(acts, outer(col, row)) ** 2)

g_col, g_row = jax.grad(factor_loss, argnums=(0, 1))(col, row)
print(jax.typeof(g_col), jax.typeof(g_row))

float32[6,5]
float32[6] float32[5]

print(jax.typeof(outer(jax.device_put(jnp.zeros(8, 'float32'), jax.P('i')),
                       row)))

import itertools

@dataclass(frozen=True)
class HiTup:
  elts: tuple

@dataclass(frozen=True)
class TupTy(HiType):
  tys: tuple  # component types: array types, or other hi types

  # lowering delegates to the component types
  def lo_ty(self):
    return [lo for ty in self.tys for lo in ty.lo_ty()]
  def lower_val(self, tup):
    return [lo for ty, elt in zip(self.tys, tup.elts)
            for lo in ty.lower_val(elt)]
  def raise_val(self, *los):
    los = iter(los)
    return HiTup(tuple(ty.raise_val(*itertools.islice(los, len(ty.lo_ty())))
                       for ty in self.tys))

  # so does the tangent type
  def to_tangent_aval(self):
    return TupTy(tuple(ty.to_tangent_aval() for ty in self.tys))

  def str_short(self, short_dtypes=False, mesh_axis_types=False):
    return ('Tup{' +
            ','.join(t.str_short(short_dtypes, mesh_axis_types)
                     for t in self.tys) + '}')
  __repr__ = str_short

register_hitype(HiTup, lambda t: TupTy(tuple(map(jax.typeof, t.elts))))

@dataclass(frozen=True)
class TupSpec(MappingSpec):
  val: tuple  # one axis entry per component

def tup_dec_rank(self, size, spec):
  return TupTy(tuple(ty.dec_rank(size, s)
                     for ty, s in zip(self.tys, spec.val)))

def tup_inc_rank(self, size, spec):
  return TupTy(tuple(ty.inc_rank(size, s)
                     for ty, s in zip(self.tys, spec.val)))

TupTy.dec_rank = tup_dec_rank
TupTy.inc_rank = tup_inc_rank

class MakeTup(VJPHiPrimitive):
  def __init__(self, elt_avals):
    self.in_avals = tuple(elt_avals)
    self.out_aval = TupTy(tuple(elt_avals))
    self.params = {}
    super().__init__()

  def expand(self, *elts):
    return HiTup(elts)

  def batch(self, axis_data, args, in_dims):
    return make_tup(*args), TupSpec(tuple(in_dims))

class GetTupElt(VJPHiPrimitive):
  def __init__(self, tup_aval, idx):
    self.in_avals = (tup_aval,)
    self.out_aval = tup_aval.tys[idx]
    self.params = dict(idx=idx)
    super().__init__()

  def expand(self, tup):
    return tup.elts[self.idx]

  def batch(self, axis_data, args, in_dims):
    tup, = args
    spec, = in_dims
    if spec is None:
      return get_tuple_element(tup, self.idx), None
    return get_tuple_element(tup, self.idx), spec.val[self.idx]

def make_tup(*elts):
  return MakeTup(map(jax.typeof, elts))(*elts)

def get_tuple_element(tup, idx):
  return GetTupElt(jax.typeof(tup), idx)(tup)

tup = make_tup(jnp.arange(3.), 5.)
print(jax.typeof(tup))
print(get_tuple_element(tup, 1))

nested = make_tup(make_tup(1., 2.), jnp.arange(2.))
print(jax.typeof(nested))
print(jax.jit(lambda t: get_tuple_element(get_tuple_element(t, 0), 1))(nested))

print(jax.typeof(make_tup(qx, 3.)))  # a quantized array element

Tup{float32[3],float32[]}
5.0
Tup{Tup{float32[],float32[]},float32[2]}
2.0
Tup{q8[2,3],float32[]}

def swap(t):
  a, b = get_tuple_element(t, 0), get_tuple_element(t, 1)
  return make_tup(b, a)

out = jax.vmap(swap, in_axes=TupSpec((0, None)), out_axes=TupSpec((None, 0)),
               axis_size=3)(tup)
print(jax.typeof(tup), '->', jax.typeof(out))

Tup{float32[3],float32[]} -> Tup{float32[],float32[3]}

print(jax.typeof(make_tup(rows, jnp.float32(1.))))

Tup{float32[8@i,3],float32[]}

What we haven’t covered

On the primitive side, there are also hooks for customizing rematerialization and dead code elimination.

As ever with hijax, tests/hijax_test.py is a good source of worked examples, and Custom derivative rules with hijax primitives covers the primitive-side API — including JVP rules, symbolic zeros, and custom linearization — in more depth.