--- jupytext: formats: md:myst text_representation: extension: .md format_name: myst format_version: 0.13 jupytext_version: 1.16.4 kernelspec: display_name: Python 3 language: python name: python3 --- (jax-301-hijax-types)= # Defining new JAX types with hijax JAX's built-in data type 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 ({ref}`jax-101-pytrees`): you bundle arrays into containers, and JAX transparently flattens the bundle into its array leaves at every boundary. But sometimes you don't want transparency. 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 rather than as a list 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 all of this with one running example: a quantized array type. We'll assume some familiarity with hijax primitives; see {ref}`jax-301-custom-derivatives` 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 `HiPrim` 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 take spec-valued `in_axes`/`out_axes` entries, plus an explicit `axis_size` when no array argument is mapped. * 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): ```{code-cell} 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 mismatching values and scales; 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 infinitesimally. 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. ```{code-cell} 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.) After registration, `jax.typeof` works on `QArray`s, 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 `HiPrim` API from {ref}`jax-301-custom-derivatives`. Each declares its input and output types, gives its implementation in `expand`, and (looking ahead to autodiff) carries a straight-through-estimator VJP rule: ```{code-cell} from jax.experimental.hijax import HiPrim class Quantize(HiPrim): 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(HiPrim): 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 `QArrayTy`s: the new type appears in primitive type signatures just like array types do. Also notice that `expand` freely constructs and inspects the `QArray` value class; primitive implementations are inside the abstraction boundary. Everything works eagerly: ```{code-cell} x = jnp.array([[1., 2., 3.], [4., -5., 6.]]) qx = quantize(x) print(qx) print(jax.typeof(qx)) print(dequantize(qx)) ``` ### 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, and 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: ```{code-cell} try: jax.jit(lambda qx: qx.qvalue)(qx) except AttributeError as e: print('AttributeError:', e) ``` Worse, calling the constructor on traced arrays doesn't fail right away. It smuggles `Tracer`s 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): ```{code-cell} 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) ``` Inside `expand` it's different: 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 too: they run eagerly, so `qx` is a genuine `QArray` holding concrete arrays. 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. The point of a quantized array type is 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: ```{code-cell} class MatmulQ(HiPrim): 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 express that rewrite once, in one place. Also notice the discipline from the previous section: `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. The backward rule passes `out_sharding` hints to its matmuls. Both are explained in the explicit-sharding section below; outside of explicit mode all these shardings are trivial and the extra code is inert.) ```{code-cell} w = jnp.arange(12., dtype='float32').reshape(3, 4) / 12. qw = quantize(w) print(matmul_q(x, qw)) print(x @ dequantize(qw)) # reference ``` ## 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: ```{code-cell} jax.jit(lambda x: dequantize(quantize(x))).trace(x).jaxpr ``` Compare to the pytree approach, where the jaxpr would inline the whole computation, and the quantized array would appear only as two unrelated intermediates, an `i8[2,3]` and an `f32[2]`, 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, as one equation with a `f32[2,3] @ q8[3,4] -> f32[2,4]` signature: ```{code-cell} jax.jit(matmul_q).trace(x, qw).jaxpr ``` And `jit` works, with quantized arrays as arguments, results, and intermediates: ```{code-cell} 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 ``` ## Autodiff and tangent types Tangent types are the main thing a pytree couldn't give us. On the type, we implemented ```python def to_tangent_aval(self): return ShapedArray(self.shape, jnp.dtype('float32'), sharding=self.sharding) ``` 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: ```{code-cell} def f(x): return jnp.sum(dequantize(quantize(x))) ``` ```{code-cell} print(jax.grad(f)(x)) ``` And differentiating with respect to a quantized array input produces a plain float array, as the tangent type dictates: ```{code-cell} def g(qx): return jnp.sum(dequantize(qx) ** 2) print(jax.grad(g)(qx)) print(jax.typeof(jax.grad(g)(qx))) ``` 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: ```{code-cell} 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)) ``` Making the tangent type an `f32` array was a *choice*, and other choices are possible. 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. ## `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 `QArray`s 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: ```{code-cell} 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: ```{code-cell} 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 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: ```{code-cell} 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 apply the unbatched operation to the stacked value, as with any 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`: ```{code-cell} 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) ``` 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: ```{code-cell} 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`, ```{code-cell} print(jax.typeof(jax.vmap(jax.jit(dequantize), in_axes=QArraySpec(), axis_size=4)(qxs))) ``` `vmap` of `grad`, and so on: ```{code-cell} 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. ```{code-cell} 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): ```{code-cell} 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]`: ```{code-cell} 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 The last piece is sharding. In JAX's explicit sharding mode (see {ref}`jax-201-sharding`), 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 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: ```{code-cell} 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))) ``` The same holds while tracing: hi types in jaxprs now display their shardings, computed by our `out_aval` rules: ```{code-cell} jax.jit(lambda x: dequantize(quantize(x))).trace(rows).jaxpr ``` 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): ```{code-cell} 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))) ``` And the contraction check fires as a type error, at trace time: ```{code-cell} 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) ``` 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 a contraction over a sharded axis should land. The right answer is the primal operand's sharding, since cotangents live where their primals live: ```{code-cell} 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)) ``` 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): ```{code-cell} 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: ```{code-cell} 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: ```{code-cell} @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) ``` And consuming one, where each device sees a per-shard `QArray` of its own rows: ```{code-cell} @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 simple: ```{code-cell} 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. ### Rank-1 arrays A rank-1 array represents an `m × n` matrix as an outer product of two vectors, `col : f32[m]` and `row : f32[n]`. The point of the representation is to never materialize the `m × n` product: ops consume the factors directly. (A general low-rank type, with `f32[m, r]` and `f32[r, n]` factors, is more of the same.) ```{code-cell} @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) ``` The type records the dense shape it represents and prints as e.g. `r1[6,5]`. Note the tangent-type choice, which differs from `QArrayTy`'s for a different reason: the sum of two rank-1 matrices is in general rank 2, so rank-1 matrices aren't closed under addition and can't serve as their own tangent space. Perturbations leave the manifold, and tangents and cotangents are dense. The sharding story is the same as before, with the field consumed in `lo_ty`: the dense row axis is carried by `col` and the dense column axis by `row`. Unlike `QArray`, whose values were only ever created inside `quantize`, here construction from factors is itself a primitive, and there's an accessor primitive too, so that *rules* (which are ordinary traced code) can get at the factors without touching container attributes: ```{code-cell} class Outer(HiPrim): 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(HiPrim): 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(HiPrim): 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) ``` (Note that `Factors` declares its output type as a *tuple* of types: `in_avals` entries and `out_aval` can be pytrees of types. It carries no autodiff rules at all, since we only apply it in forward passes; rules are only needed for the transformations you actually use. Note also that `MatmulR1`'s backward rule computes the dense operand's cotangent as an outer product, exploiting the representation the same way `expand` does.) ```{code-cell} 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)))) ``` Jaxprs, `jit`, and autodiff all work as before. The gradient with respect to a rank-1 operand is dense, as the tangent type dictates, while gradients chained through the constructor land back on the factors: ```{code-cell} jax.jit(matmul_r1).trace(acts, r1).jaxpr ``` ```{code-cell} 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)) ``` And the typing rules propagate shardings, just like `QArrayTy`'s: ```{code-cell} print(jax.typeof(outer(jax.device_put(jnp.zeros(8, 'float32'), jax.P('i')), row))) ``` We stopped here, but `vmap`, `scan`, and `shard_map` support would follow the same recipes as in the sections above: `dec_rank`/`inc_rank` and a mapping spec, `leading_axis_spec`, and `shard`/`unshard` with an `HiPspec` partition spec type. ### Tuples Our last example is a generic container: a tuple whose elements are any JAX values. Where `QArrayTy` and `Rank1Ty` had a fixed component structure, `TupTy` is parameterized by its component *types*, and every method delegates to them: ```{code-cell} 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)))) ``` Two things follow from the delegation. First, elements can themselves be hi types (tuples of tuples, or a tuple holding a `QArray`), since `lo_ty` and friends recurse. Second, there's no `sharding` field this time: the stored component types carry their own shardings, a third way of handling sharding-in-types alongside `QArrayTy`'s single field and `Rank1Ty`'s dense-shape spec. The constructor and accessor primitives are familiar from the rank-1 example. One new ingredient: the element index is a *static parameter*, passed via `params` and available as `self.idx`. And since each element can be batched along its own axis (or not at all), the mapping spec carries one entry per component: ```{code-cell} @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(HiPrim): 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(HiPrim): 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) ``` (Note `GetTupElt`'s batch rule handles unbatched inputs: hi primitive batch rules are invoked even when no argument is batched, so `in_dims` entries can be `None`.) Tuples work, they nest, and they hold other hi types: ```{code-cell} 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 ``` Under `vmap`, the per-component mapping spec gives each element its own `in_axes`/`out_axes` entry, visible in the types. Here the first element is mapped on the way in, and only the second on the way out: ```{code-cell} 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)) ``` And since the component types carry their own shardings, sharding-in-types needs nothing extra: ```{code-cell} print(jax.typeof(make_tup(rows, jnp.float32(1.)))) ``` For autodiff rules on these primitives, and `scan` and `shard_map` support (via a per-component `HiPspec`), see the `TupTy` example in `tests/hijax_test.py`. ## 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 {ref}`jax-301-custom-derivatives` covers the primitive-side API (including JVP rules, symbolic zeros, and custom linearization) in more depth.