Defining new JAX types with hijax — JAX documentation

Skip to main content

Back to top

Ctrl+K

Search<br>Ctrl+K

.ipynb

.pdf

Defining new JAX types with hijax

Contents

Defining new JAX types with hijax#

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

But sometimes transparency is exactly what you don’t want. Some data is best<br>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<br>spray of array leaves;

it has internal invariants, so users should only produce and consume it<br>through a fixed set of operations, rather than by freely constructing or<br>pattern-matching its components;

its tangent type may differ from its primal structure, so that<br>derivatives with respect to it aren’t just “the same pytree, but for<br>tangents”;

it may have its own notion of batching under vmap;

it can carry sharding information in the type, participating in JAX’s<br>explicit sharding mode.

Hijax types (or “hi types”) provide this. You subclass HiType to define<br>the type, register a Python class as carrying values of that type, and write<br>hijax primitives whose input and output types mention the new type. This<br>document walks through the whole story with one running example: a<br>quantized array type.

We’ll assume some familiarity with hijax primitives; see<br>Custom derivative rules with hijax primitives for an introduction to them. Like everything<br>hijax, this is experimental: expect imports from jax.experimental.hijax,<br>and expect the APIs to evolve.

TL;DR#

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

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

For autodiff, implement to_tangent_aval on the type, and VJP/JVP rules<br>on the primitives.

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

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

Example: quantized arrays#

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

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

from dataclasses import dataclass

import jax<br>import jax.numpy as jnp

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

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

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

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

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

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

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

We also give the type a sharding field, recording how values...