Defining new Jax types with hijax

Defining new JAX types with hijax

JAX's built-in currency is the array: functions you transform take arrays in and produce arrays out, and every intermediate the tracing machinery sees has an array type like f32[3,4]. When you want to work with aggregate data, the usual tool is a pytree: you bundle arrays into containers, and JAX transparently flattens the bundle into its array leaves at every boundary.

But sometimes transparency is exactly what you don't want. Some data is best modeled as a new type, with its own identity:

• It should appear in jaxprs as a single value of a single type, not as a spray of array leaves
• It has internal invariants, so users should only produce and consume it through a fixed set of operations, rather than by freely constructing or pattern-matching its components
• Its tangent type may differ from its primal structure, so that derivatives with respect to it aren't just "the same pytree, but for tangents"
• It may have its own notion of batching under vmap
• It can carry sharding information in the type, participating in JAX's explicit sharding mode

Hijax types (or "hi types") provide this. You subclass HiType to define the type, register a Python class as carrying values of that type, and write hijax primitives whose input and output types mention the new type. This document walks through the whole story with one running example: a quantized array type. We'll assume some familiarity with hijax primitives; see Custom derivative rules with hijax primitives for an introduction to them. Like everything hijax, this is experimental: expect imports from jax.experimental.hijax, and expect the APIs to evolve.

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.      ]]

Primitives outside, container inside

The expand methods above construct QArray values and read their attributes directly. It's important that this kind of direct container manipulation happen only in expand (and in the type's own methods, like lower_val and raise_val). Everywhere else - in particular, in any function you might jit, differentiate, or vmap - hi values should be produced and consumed only by applying primitives.

The reason is what traced code actually sees. Under a trace, a quantized array is not a QArray instance: it's a Tracer of type q8[...]. So reading an attribute in traced code fails:

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

AttributeError: DynamicJaxprTracer has no attribute qvalue

And, worse, calling the constructor on traced arrays doesn't fail right away. It smuggles Tracers inside a container that JAX treats as one opaque concrete value, and the mistake surfaces later, as a confusing error far from its cause (here a missing constant handler; under grad it's a leaked-tracer error):

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 '__main__.QArray'>

In expand it's a different story: by the time expand runs, JAX has committed to implementing the primitive in terms of the type's lojax components, and its QArray arguments are genuine QArray instances (holding lojax values - possibly traced ones). There, manipulating the value as a plain container is exactly right. (The top-level peeks at attributes like qx.qvalue elsewhere in this document are fine for the same reason eager expand is fine: they run eagerly, on concrete values. But inside any function that might get traced, stick to primitives.)

A more realistic op: dense × quantized matmul

Conversion ops alone make for a thin API: in practice, a quantized array type earns its keep in ops that consume the type directly. The classic example is an inference-style matmul, where the activations x are an ordinary dense f32 array and the weights are quantized:

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)

A few things to notice. The type signature mixes array types and hi types freely: in_avals is a (ShapedArray, QArrayTy) pair, checked at construction time so that a shape mismatch fails immediately, with both operand types pretty-printed. And expand exploits the representation: because the scales apply per-row along the contraction axis, they can be folded into the dense operand, so the heavy matmul runs directly against the int8 payload rather than a dequantized copy. Owning the op as a single primitive lets us state that rewriting once, in one place.

Also notice the discipline from the previous section in action: expand reads qw.scale and qw.qvalue as container attributes, while the VJP rules - which are ordinary traced code - go through the dequantize primitive instead. (The typing rule also computes an output sharding - output rows partitioned like x's, output columns like qw's, with the contracted axes required to be unsharded - and the backward rule passes out_sharding hints to its matmuls. Both are explained in the explicit-sharding section.)