{ "cells": [ { "cell_type": "raw", "id": "02454120", "metadata": {}, "source": [ "---\n", "Copyright 2026 The JAX Authors.\n", "\n", "Licensed under the Apache License, Version 2.0 (the \"License\");\n", "you may not use this file except in compliance with the License.\n", "You may obtain a copy of the License at\n", "\n", " https://www.apache.org/licenses/LICENSE-2.0\n", "\n", "Unless required by applicable law or agreed to in writing, software\n", "distributed under the License is distributed on an \"AS IS\" BASIS,\n", "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "See the License for the specific language governing permissions and\n", "limitations under the License.\n", "\n", "---" ] }, { "cell_type": "markdown", "id": "d3a34514", "metadata": {}, "source": [ "(hijax-types)=\n", "# Defining new JAX types with hijax\n", "\n", "JAX's built-in currency is the array: functions you transform take arrays in\n", "and produce arrays out, and every intermediate the tracing machinery sees has\n", "an array type like `f32[3,4]`. When you want to work with aggregate data, the\n", "usual tool is a\n", "[pytree](https://docs.jax.dev/en/latest/pytrees.html): you\n", "bundle arrays into containers, and JAX transparently flattens the bundle\n", "into its array leaves at every boundary.\n", "\n", "But sometimes transparency is exactly what you don't want. Some data is best\n", "modeled as a new *type*, with its own identity:\n", "\n", "* it should appear in jaxprs as a single value of a single type, not as a\n", " spray of array leaves;\n", "* it has internal invariants, so users should only produce and consume it\n", " through a fixed set of operations, rather than by freely constructing or\n", " pattern-matching its components;\n", "* its *tangent type* may differ from its primal structure, so that\n", " derivatives with respect to it aren't just \"the same pytree, but for\n", " tangents\";\n", "* it may have its own notion of batching under `vmap`;\n", "* it can carry sharding information in the type, participating in JAX's\n", " explicit sharding mode.\n", "\n", "Hijax types (or \"hi types\") provide this. You subclass `HiType` to define\n", "the type, register a Python class as carrying values of that type, and write\n", "hijax primitives whose input and output types mention the new type. This\n", "document walks through the whole story with one running example: a\n", "quantized array type.\n", "\n", "We'll assume some familiarity with hijax primitives; see\n", "{ref}`hijax-custom-derivatives` for an introduction to them. Like everything\n", "hijax, this is experimental: expect imports from `jax.experimental.hijax`,\n", "and expect the APIs to evolve.\n", "\n", "### TL;DR\n", "\n", "* Subclass `HiType` and implement `lo_ty`, `lower_val`, and `raise_val` to\n", " say how the type and its values lower to ordinary (\"lojax\") arrays, then\n", " call `register_hitype` to associate your value class with your type.\n", "* Write `HiPrim` subclasses whose `in_avals`/`out_aval` mention the\n", " new type; these are the only way values of the type get produced and\n", " consumed.\n", "* For autodiff, implement `to_tangent_aval` on the type, and VJP/JVP rules\n", " on the primitives.\n", "* For `vmap`, implement `dec_rank` and `inc_rank` on the type along with a\n", " `MappingSpec` subclass of your own design, and `batch` rules on the\n", " primitives. Mapped-over hi type arguments require an explicit `axis_size`\n", " and spec-valued `in_axes`/`out_axes` entries.\n", "* For sharding in types (explicit mode), record sharding data on your type\n", " (e.g. a `NamedSharding` field), consume it in `lo_ty`, and propagate it\n", " in your primitives' typing rules.\n", "\n", "## Example: quantized arrays\n", "\n", "Say we want to work with arrays quantized to `int8`. A quantized array is\n", "really a pair of arrays: the `int8` values, and a floating point scale\n", "shared by each row (that is, we quantize along the last axis, one scale per\n", "row, as in common per-row/per-channel quantization schemes):" ] }, { "cell_type": "code", "execution_count": null, "id": "a6a0201f", "metadata": {}, "outputs": [], "source": [ "import os\n", "os.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'\n", "# (8 CPU devices, for the sharding sections at the end)\n", "\n", "from dataclasses import dataclass\n", "\n", "import jax\n", "import jax.numpy as jnp\n", "\n", "@dataclass(frozen=True)\n", "class QArray:\n", " qvalue: jax.Array # int8[*leading, n]\n", " scale: jax.Array # f32[*leading]" ] }, { "cell_type": "markdown", "id": "f4d61d4b", "metadata": {}, "source": [ "We could register `QArray` as a pytree and be done. But consider what we'd\n", "give up:\n", "\n", "* **Invariants.** The two components are coupled: `scale` must have the\n", " shape of `qvalue` minus its last axis, and `qvalue` is only meaningful\n", " together with its `scale`. As a pytree, nothing stops code from crossing\n", " the streams; under transformations, JAX itself sees only independent\n", " leaves.\n", "* **Types in jaxprs.** As a pytree, a quantized array appears in traced\n", " code as two unrelated array values. We'd rather see one value, of one\n", " type, so jaxprs say what they mean.\n", "* **Tangents.** A quantized array's values live on a discrete grid, so it\n", " makes no sense to perturb them along the grid. But a pytree's tangent\n", " type is forced to be the pytree of its leaves' tangent types — and the\n", " tangent type of an integer array like `qvalue` is a `float0` array,\n", " which can only carry a trivial payload. So as a pytree, a quantized\n", " array would admit no useful perturbations at all. What we want is to\n", " choose a tangent type for the quantized array as a whole, such as the\n", " *continuous* `f32` arrays that the quantized values approximate.\n", "\n", "So instead we'll make `QArray` a hijax type.\n", "\n", "## The type\n", "\n", "A hijax type is a subclass of `HiType`. The required core is small:\n", "\n", "* `lo_ty` says which lojax (array) types make up the type;\n", "* `lower_val` and `raise_val` convert values to and from that list of\n", " arrays;\n", "* the type must be hashable and comparable for equality (a frozen dataclass\n", " gives us both).\n", "\n", "This is like the pytree flatten/unflatten interface, but it lives at the\n", "level of *types*: given only the type, JAX can compute the lowered types,\n", "without needing a value in hand.\n", "\n", "We also give the type a `sharding` field, recording how values are\n", "partitioned across devices; it can be ignored until the sharding sections\n", "near the end of this document. We reuse JAX's `NamedSharding` to describe\n", "the partitioning of the `qvalue` component, and derive `scale`'s\n", "partitioning from it by dropping the last axis. Nothing about the field is\n", "special to JAX, which never interprets it: only our own methods consume\n", "it, chiefly `lo_ty`, which stamps the component types with their\n", "shardings. You're free to track sharding information on your type with an\n", "object of your own design instead, so long as you consume it the same\n", "way." ] }, { "cell_type": "code", "execution_count": null, "id": "25ca92c8", "metadata": {}, "outputs": [], "source": [ "from jax.experimental.hijax import HiType, ShapedArray, register_hitype\n", "from jax.sharding import NamedSharding\n", "\n", "@dataclass(frozen=True)\n", "class QArrayTy(HiType):\n", " shape: tuple[int, ...]\n", " sharding: NamedSharding # qvalue's sharding; scale's is derived from it\n", "\n", " # lowering: which array types make up this type, and how values convert\n", " def lo_ty(self):\n", " scale_sharding = self.sharding.update(spec=jax.P(*self.sharding.spec[:-1]))\n", " return [ShapedArray(self.shape, jnp.dtype('int8'),\n", " sharding=self.sharding),\n", " ShapedArray(self.shape[:-1], jnp.dtype('float32'),\n", " sharding=scale_sharding)]\n", " def lower_val(self, q):\n", " return [q.qvalue, q.scale]\n", " def raise_val(self, qvalue, scale):\n", " return QArray(qvalue, scale)\n", "\n", " # autodiff: tangents of quantized arrays are plain float arrays (see below)\n", " def to_tangent_aval(self):\n", " return ShapedArray(self.shape, jnp.dtype('float32'),\n", " sharding=self.sharding)\n", "\n", " # printing, e.g. in jaxprs\n", " def str_short(self, short_dtypes=False, mesh_axis_types=False):\n", " dims = [str(d) if p is None else f'{d}@{p}'\n", " for d, p in zip(self.shape, self.sharding.spec)]\n", " return f'q8[{\",\".join(dims)}]'\n", " __repr__ = str_short\n", "\n", "register_hitype(QArray, lambda q: QArrayTy(q.qvalue.shape,\n", " jax.typeof(q.qvalue).sharding))" ] }, { "cell_type": "markdown", "id": "c2cb7768", "metadata": {}, "source": [ "The `register_hitype` call associates the value class with the type: its\n", "second argument computes the type of any given value, analogous to how\n", "`jax.typeof` maps an array to its `ShapedArray` type. (Ours reads both the\n", "shape and the sharding off the `qvalue` component — every array carries a\n", "sharding, trivial when no mesh is in play.) Indeed after registration,\n", "`jax.typeof` works on `QArray`s, and JAX transformations accept them\n", "anywhere a value is expected.\n", "\n", "## The primitives\n", "\n", "With a pytree, users construct and take apart values freely. With a hijax\n", "type, values are produced and consumed only by hijax primitives whose\n", "declared types mention the new type. That's where invariants get enforced:\n", "if every primitive preserves them, they always hold.\n", "\n", "Our two primitives are `quantize` and `dequantize`, written with the\n", "`HiPrim` API from {ref}`hijax-custom-derivatives`. Each declares\n", "its input and output types, gives its implementation in `expand`, and\n", "(looking ahead to autodiff) carries a straight-through-estimator VJP rule:" ] }, { "cell_type": "code", "execution_count": null, "id": "645f5520", "metadata": {}, "outputs": [], "source": [ "from jax.experimental.hijax import HiPrim\n", "\n", "class Quantize(HiPrim):\n", " def __init__(self, x_aval):\n", " if x_aval.dtype != jnp.dtype('float32'): raise TypeError(x_aval.dtype)\n", " self.in_avals = (x_aval,)\n", " self.out_aval = QArrayTy(x_aval.shape, x_aval.sharding)\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x):\n", " scale = jnp.max(jnp.abs(x), axis=-1) / 127.\n", " qvalue = jnp.round(x / scale[..., None]).astype(jnp.int8)\n", " return QArray(qvalue, scale)\n", "\n", " # straight-through estimator: differentiate as if it's the identity\n", " def vjp_fwd(self, nzs_in, x):\n", " return self(x), None\n", "\n", " def vjp_bwd_retval(self, _res, g):\n", " return (g,)\n", "\n", "class Dequantize(HiPrim):\n", " def __init__(self, q_aval):\n", " self.in_avals = (q_aval,)\n", " self.out_aval = ShapedArray(q_aval.shape, jnp.dtype('float32'),\n", " sharding=q_aval.sharding)\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, qx):\n", " return qx.qvalue.astype('float32') * qx.scale[..., None]\n", "\n", " def vjp_fwd(self, nzs_in, qx):\n", " return self(qx), None\n", "\n", " def vjp_bwd_retval(self, _res, g):\n", " return (g,)\n", "\n", "def quantize(x):\n", " return Quantize(jax.typeof(x))(x)\n", "\n", "def dequantize(qx):\n", " return Dequantize(jax.typeof(qx))(qx)" ] }, { "cell_type": "markdown", "id": "6aefa0bb", "metadata": {}, "source": [ "Notice that `Quantize`'s `out_aval` and `Dequantize`'s `in_avals` are\n", "`QArrayTy`s: the new type appears in primitive type signatures just like\n", "array types do. Also notice `expand` freely constructs and inspects the\n", "`QArray` value class; primitive implementations are inside the abstraction\n", "boundary.\n", "\n", "Everything works eagerly:" ] }, { "cell_type": "code", "execution_count": null, "id": "0fedd0ea", "metadata": {}, "outputs": [], "source": [ "x = jnp.array([[1., 2., 3.],\n", " [4., -5., 6.]])\n", "\n", "qx = quantize(x)\n", "print(qx)\n", "print(jax.typeof(qx))\n", "print(dequantize(qx))" ] }, { "cell_type": "markdown", "id": "49ccc6a2", "metadata": {}, "source": [ "### Primitives outside, container inside\n", "\n", "The `expand` methods above construct `QArray` values and read their\n", "attributes directly. It's important that this kind of direct container\n", "manipulation happen *only* in `expand` (and in the type's own methods,\n", "like `lower_val` and `raise_val`). Everywhere else — in particular, in\n", "any function you might `jit`, differentiate, or `vmap` — hi values should\n", "be produced and consumed only by applying primitives.\n", "\n", "The reason is what traced code actually sees. Under a trace, a quantized\n", "array is not a `QArray` instance: it's a `Tracer` of type `q8[...]`. So\n", "reading an attribute in traced code fails —" ] }, { "cell_type": "code", "execution_count": null, "id": "de0d922c", "metadata": {}, "outputs": [], "source": [ "try:\n", " jax.jit(lambda qx: qx.qvalue)(qx)\n", "except AttributeError as e:\n", " print('AttributeError:', e)" ] }, { "cell_type": "markdown", "id": "47ef5033", "metadata": {}, "source": [ "and, worse, calling the constructor on traced arrays doesn't fail right\n", "away. It smuggles `Tracer`s inside a container that JAX treats as one\n", "opaque concrete value, and the mistake surfaces later, as a confusing\n", "error far from its cause (here a missing constant handler; under `grad`\n", "it's a leaked-tracer error):" ] }, { "cell_type": "code", "execution_count": null, "id": "a8c7ed79", "metadata": {}, "outputs": [], "source": [ "def bad_quantize(x):\n", " scale = jnp.max(jnp.abs(x), axis=-1) / 127.\n", " return QArray(jnp.round(x / scale[..., None]).astype('int8'), scale)\n", "\n", "try:\n", " jax.jit(bad_quantize)(x)\n", "except TypeError as e:\n", " print('TypeError:', e)" ] }, { "cell_type": "markdown", "id": "645c78bd", "metadata": {}, "source": [ "In `expand` it's a different story: by the time `expand` runs, JAX has\n", "committed to implementing the primitive in terms of the type's lojax\n", "components, and its `QArray` arguments are genuine `QArray` instances\n", "(holding lojax values — possibly traced ones). There, manipulating the\n", "value as a plain container is exactly right.\n", "\n", "(The top-level peeks at attributes like `qx.qvalue` elsewhere in this\n", "document are fine for the same reason eager `expand` is fine: they run\n", "eagerly, on concrete values. But inside any function that might get\n", "traced, stick to primitives.)\n", "\n", "### A more realistic op: dense × quantized matmul\n", "\n", "Conversion ops alone make for a thin API: in practice, a quantized array\n", "type earns its keep in ops that consume the type directly. The classic\n", "example is an inference-style matmul, where the activations `x` are an\n", "ordinary dense `f32` array and the weights are quantized:" ] }, { "cell_type": "code", "execution_count": null, "id": "18416d37", "metadata": {}, "outputs": [], "source": [ "class MatmulQ(HiPrim):\n", " def __init__(self, x_aval, q_aval):\n", " if not (isinstance(q_aval, QArrayTy) and len(x_aval.shape) == 2 and\n", " len(q_aval.shape) == 2 and x_aval.shape[1] == q_aval.shape[0]):\n", " raise TypeError(f'bad matmul_q operand types: {x_aval} @ {q_aval}')\n", " self.in_avals = (x_aval, q_aval)\n", " x_spec, q_spec = x_aval.sharding.spec, q_aval.sharding.spec\n", " if x_spec[1] is not None or q_spec[0] is not None:\n", " raise TypeError('matmul_q requires unsharded contraction axes, got '\n", " f'{x_aval} @ {q_aval}')\n", " out_sharding = x_aval.sharding.update(spec=jax.P(x_spec[0], q_spec[1]))\n", " self.out_aval = ShapedArray((x_aval.shape[0], q_aval.shape[1]),\n", " jnp.dtype('float32'), sharding=out_sharding)\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x, qw):\n", " # fold the per-row scales into the dense operand, then apply one matmul\n", " # directly against the int8 payload\n", " return (x * qw.scale) @ qw.qvalue.astype(jnp.float32)\n", "\n", " def vjp_fwd(self, nzs_in, x, qw):\n", " return self(x, qw), (x, qw)\n", "\n", " def vjp_bwd_retval(self, res, g):\n", " x, qw = res\n", " w = dequantize(qw) # rules are traced code: use primitives here\n", " # cotangents live where the primals live, so use the primal operands'\n", " # shardings to disambiguate the (possibly sharded) contractions\n", " return (jnp.matmul(g, w.T, out_sharding=jax.typeof(x).sharding),\n", " jnp.matmul(x.T, g, out_sharding=jax.typeof(w).sharding))\n", "\n", "def matmul_q(x, qw):\n", " return MatmulQ(jax.typeof(x), jax.typeof(qw))(x, qw)" ] }, { "cell_type": "markdown", "id": "4dff6f01", "metadata": {}, "source": [ "A few things to notice. The type signature mixes array types and hi types\n", "freely: `in_avals` is a `(ShapedArray, QArrayTy)` pair, checked at\n", "construction time so that a shape mismatch fails immediately, with both\n", "operand types pretty-printed. And `expand` exploits the representation:\n", "because the scales apply per-row along the contraction axis, they can be\n", "folded into the dense operand, so the heavy matmul runs directly against\n", "the `int8` payload rather than a dequantized copy. Owning the op as a\n", "single primitive lets us state that rewriting once, in one place.\n", "\n", "Also notice the discipline from the previous section in action: `expand`\n", "reads `qw.scale` and `qw.qvalue` as container attributes, while the VJP\n", "rules — which are ordinary traced code — go through the `dequantize`\n", "primitive instead.\n", "\n", "(The typing rule also computes an output *sharding* — output rows\n", "partitioned like `x`'s, output columns like `qw`'s, with the contracted\n", "axes required to be unsharded — and the backward rule passes\n", "`out_sharding` hints to its matmuls. Both are explained in the\n", "explicit-sharding section below; outside of explicit mode all these\n", "shardings are trivial and the extra code is inert.)" ] }, { "cell_type": "code", "execution_count": null, "id": "e238c0f9", "metadata": {}, "outputs": [], "source": [ "w = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.\n", "qw = quantize(w)\n", "\n", "print(matmul_q(x, qw))\n", "print(x @ dequantize(qw)) # reference" ] }, { "cell_type": "markdown", "id": "ee8a9685", "metadata": {}, "source": [ "## Hi types in jaxprs\n", "\n", "When we trace, the quantized array appears as a single value of type\n", "`q8[2,3]`, produced by one equation and consumed by another:" ] }, { "cell_type": "code", "execution_count": null, "id": "dddcbb16", "metadata": {}, "outputs": [], "source": [ "jax.jit(lambda x: dequantize(quantize(x))).trace(x).jaxpr" ] }, { "cell_type": "markdown", "id": "a04546c6", "metadata": {}, "source": [ "Compare to the pytree approach, where the same computation would show four\n", "array-typed intermediates with no indication that they pair up. The hi type\n", "only disappears at lowering time, when `expand` is traced and each\n", "`q8[...]`-typed value is expanded into the array components given by\n", "`lo_ty`.\n", "\n", "Ops with mixed operand kinds read just as directly — one equation with a\n", "`f32[2,3] @ q8[3,4] -> f32[2,4]` signature:" ] }, { "cell_type": "code", "execution_count": null, "id": "cba8a58c", "metadata": {}, "outputs": [], "source": [ "jax.jit(matmul_q).trace(x, qw).jaxpr\n", "#\n", "# `jit` works, with quantized arrays as arguments, results, and\n", "# intermediates:" ] }, { "cell_type": "code", "execution_count": null, "id": "bb8b2b69", "metadata": {}, "outputs": [], "source": [ "print(jax.jit(lambda x: dequantize(quantize(x)))(x)) # QArray internal\n", "\n", "qx2 = jax.jit(quantize)(x) # QArray result\n", "print(jax.typeof(qx2))\n", "\n", "print(jax.jit(dequantize)(qx2)) # QArray argument" ] }, { "cell_type": "markdown", "id": "b95205e3", "metadata": {}, "source": [ "## Autodiff and tangent types\n", "\n", "Here's where hi types earn their keep. On the type, we implemented\n", "\n", "```python\n", " def to_tangent_aval(self):\n", " return ShapedArray(self.shape, jnp.dtype('float32'))\n", "```\n", "\n", "which says: the tangent type of a quantized array is a plain `f32` array.\n", "No pytree can express this: a pytree's tangent type is always the pytree\n", "of its leaves' tangent types, and for the `int8` leaf `qvalue` that means\n", "a trivial `float0` tangent.\n", "\n", "Together with the straight-through VJP rules on the primitives, gradients\n", "flow through quantization as if it were the identity:" ] }, { "cell_type": "code", "execution_count": null, "id": "4595e881", "metadata": {}, "outputs": [], "source": [ "def f(x):\n", " return jnp.sum(dequantize(quantize(x)))" ] }, { "cell_type": "code", "execution_count": null, "id": "3a2203a8", "metadata": {}, "outputs": [], "source": [ "print(jax.grad(f)(x))" ] }, { "cell_type": "markdown", "id": "7936d5a9", "metadata": {}, "source": [ "And differentiating with respect to a quantized array input produces a\n", "plain float array, as the tangent type dictates:" ] }, { "cell_type": "code", "execution_count": null, "id": "6e170986", "metadata": {}, "outputs": [], "source": [ "def g(qx):\n", " return jnp.sum(dequantize(qx) ** 2)\n", "\n", "print(jax.grad(g)(qx))\n", "print(jax.typeof(jax.grad(g)(qx)))" ] }, { "cell_type": "markdown", "id": "a77c3e02", "metadata": {}, "source": [ "The same goes for ops with mixed operand kinds, like the quantized\n", "matmul: differentiating a loss with respect to both operands gives an\n", "`f32` gradient for the dense activations *and* an `f32` gradient for the\n", "quantized weights:" ] }, { "cell_type": "code", "execution_count": null, "id": "ff1d870e", "metadata": {}, "outputs": [], "source": [ "def loss(x, qw):\n", " return jnp.sum(matmul_q(x, qw) ** 2)\n", "\n", "grad_x, grad_qw = jax.grad(loss, argnums=(0, 1))(x, qw)\n", "print(jax.typeof(grad_x), jax.typeof(grad_qw))" ] }, { "cell_type": "markdown", "id": "6ea2bf8c", "metadata": {}, "source": [ "Notice that making the tangent type an `f32` array was a *choice*, and\n", "there's a real design space here. We could instead have made the tangent\n", "type of `QArrayTy` be `QArrayTy` itself, so that tangents and cotangents\n", "are quantized too — a different tradeoff, sensible for different\n", "applications. (For that choice, since the tangent type is then a hi type,\n", "we'd also implement `vspace_zero` and `vspace_add` on it so autodiff can\n", "instantiate and accumulate cotangents.) This flexibility is why hi types\n", "are a user extension point: for each piece of JAX — tracing, lowering,\n", "autodiff, and batching — you set up how your type participates, however\n", "your situation needs.\n", "\n", "## `vmap` and mapping specs\n", "\n", "What does it mean to map over a quantized array? For arrays, `vmap`'s\n", "`in_axes` and `out_axes` are axis indices, and JAX can infer the mapped\n", "axis size from the argument's shape. For a general hi type, JAX doesn't\n", "guess: *you* define a \"mapping spec\" type that says how values of your type\n", "are mapped, users pass instances of it as `in_axes`/`out_axes` entries, and\n", "they pass `axis_size` explicitly when it can't be inferred from an array\n", "argument.\n", "\n", "For our quantized arrays, thanks to the per-row scales, a batch of\n", "`QArray`s is just a bigger `QArray`: stacking `n` quantized arrays of type\n", "`q8[2,3]` along a new leading axis gives a `q8[n,2,3]`, with `qvalue` of\n", "shape `(n, 2, 3)` and `scale` of shape `(n, 2)`. So the only mapping notion\n", "we need is \"the leading axis,\" and our spec type doesn't need to carry any\n", "data at all:" ] }, { "cell_type": "code", "execution_count": null, "id": "08b53a74", "metadata": {}, "outputs": [], "source": [ "from jax.experimental.hijax import MappingSpec\n", "\n", "@dataclass(frozen=True)\n", "class QArraySpec(MappingSpec):\n", " pass # QArrays are only mapped along their leading axis" ] }, { "cell_type": "markdown", "id": "d62ddc51", "metadata": {}, "source": [ "(Specs can be as rich as your type demands. A tuple-like hi type might use\n", "a spec carrying one axis per component; see the tuple example at the end\n", "of this document.)\n", "\n", "On the type, we implement `dec_rank` and `inc_rank`, the hi type analogues\n", "of \"remove the mapped axis\" and \"add the mapped axis.\" They take the axis\n", "size and a spec, and return the element type and the batched type,\n", "respectively:" ] }, { "cell_type": "code", "execution_count": null, "id": "3abd307e", "metadata": {}, "outputs": [], "source": [ "def qarray_dec_rank(self, size, spec):\n", " assert isinstance(spec, QArraySpec) and self.shape[0] == size\n", " return QArrayTy(self.shape[1:],\n", " self.sharding.update(spec=jax.P(*self.sharding.spec[1:])))\n", "\n", "def qarray_inc_rank(self, size, spec):\n", " assert isinstance(spec, QArraySpec)\n", " return QArrayTy((size, *self.shape),\n", " self.sharding.update(spec=jax.P(None, *self.sharding.spec)))\n", "\n", "QArrayTy.dec_rank = qarray_dec_rank\n", "QArrayTy.inc_rank = qarray_inc_rank" ] }, { "cell_type": "markdown", "id": "34931284", "metadata": {}, "source": [ "(We're attaching methods to the class as we go, notebook-style; in real\n", "code these would just be more methods in the `class QArrayTy` definition.)\n", "\n", "On the primitives, we implement `batch` rules. A `batch` rule receives the\n", "batched arguments along with their mapping specs (`None` for unbatched\n", "arguments, an integer axis for batched array arguments, and a spec instance\n", "for batched hi type arguments), and returns the batched result along with\n", "its mapping spec. Note that a rule should be prepared for any combination\n", "of batched and unbatched arguments:" ] }, { "cell_type": "code", "execution_count": null, "id": "934ebaea", "metadata": {}, "outputs": [], "source": [ "def quantize_batch(self, axis_data, args, in_dims):\n", " x, = args\n", " d, = in_dims\n", " if d is None:\n", " return quantize(x), None\n", " x = jnp.moveaxis(x, d, 0)\n", " return quantize(x), QArraySpec()\n", "Quantize.batch = quantize_batch\n", "\n", "def dequantize_batch(self, axis_data, args, in_dims):\n", " qx, = args\n", " d, = in_dims\n", " if d is None:\n", " return dequantize(qx), None\n", " assert isinstance(d, QArraySpec)\n", " return dequantize(qx), 0\n", "Dequantize.batch = dequantize_batch" ] }, { "cell_type": "markdown", "id": "a256ff90", "metadata": {}, "source": [ "Because per-row quantization applies at any rank, both rules can just apply\n", "the unbatched operation to the stacked value — the hallmark of a type whose\n", "batches are values of the same type family.\n", "\n", "Now we can `vmap`. Mapping *to* a quantized array output, the axis size is\n", "inferred from the array argument as usual, and we pass a spec for\n", "`out_axes`:" ] }, { "cell_type": "code", "execution_count": null, "id": "a4ec65e2", "metadata": {}, "outputs": [], "source": [ "xs = jnp.arange(24., dtype='float32').reshape(4, 2, 3)\n", "\n", "qxs = jax.vmap(quantize, out_axes=QArraySpec())(xs)\n", "print(jax.typeof(qxs))\n", "print(qxs.qvalue.shape, qxs.scale.shape)" ] }, { "cell_type": "markdown", "id": "4fe7acd4", "metadata": {}, "source": [ "Mapping *over* a quantized array input, we pass a spec for `in_axes` — and\n", "since there's no array argument to infer the axis size from, we must pass\n", "`axis_size` explicitly:" ] }, { "cell_type": "code", "execution_count": null, "id": "8d347d8f", "metadata": {}, "outputs": [], "source": [ "xs_roundtrip = jax.vmap(dequantize, in_axes=QArraySpec(), axis_size=4)(qxs)\n", "print(jax.typeof(xs_roundtrip))" ] }, { "cell_type": "markdown", "id": "5918261f", "metadata": {}, "source": [ "All the usual compositions work — `vmap` of `jit`," ] }, { "cell_type": "code", "execution_count": null, "id": "b687b6b3", "metadata": {}, "outputs": [], "source": [ "print(jax.typeof(jax.vmap(jax.jit(dequantize), in_axes=QArraySpec(),\n", " axis_size=4)(qxs)))" ] }, { "cell_type": "markdown", "id": "81f832d9", "metadata": {}, "source": [ "`vmap` of `grad`, and so on:" ] }, { "cell_type": "code", "execution_count": null, "id": "c2eb857d", "metadata": {}, "outputs": [], "source": [ "def norm_quantized(x):\n", " return jnp.sum(dequantize(quantize(x)) ** 2)\n", "\n", "print(jax.vmap(jax.grad(norm_quantized))(xs).shape)" ] }, { "cell_type": "markdown", "id": "df974fde", "metadata": {}, "source": [ "## `scan` and the leading axis\n", "\n", "`jax.lax.scan` can loop over a stacked hi value, consuming one slice per\n", "step — and it can carry and produce hi values too. Where `vmap` asks the\n", "*user* for a mapping spec, `scan` always walks the leading axis, so it\n", "instead asks the *type*: the one extra method to implement is\n", "`leading_axis_spec`, which returns the mapping spec describing your type's\n", "leading axis. The `dec_rank` and `inc_rank` methods from the `vmap`\n", "section do the rest of the work." ] }, { "cell_type": "code", "execution_count": null, "id": "dc33fb02", "metadata": {}, "outputs": [], "source": [ "def qarray_leading_axis_spec(self):\n", " return QArraySpec()\n", "\n", "QArrayTy.leading_axis_spec = qarray_leading_axis_spec" ] }, { "cell_type": "markdown", "id": "aec28628", "metadata": {}, "source": [ "Scanning over the stack of quantized arrays from the previous section, the\n", "body sees one `q8[2,3]` per step. As with `vmap`'s `axis_size`, when all\n", "the scanned-over values are hi types there's no leading-axis size to\n", "infer, so we pass `length` explicitly (if any scanned-over value is an\n", "array, `scan` infers the length from it and `length` can be omitted):" ] }, { "cell_type": "code", "execution_count": null, "id": "ea5c42f8", "metadata": {}, "outputs": [], "source": [ "def sum_dequantized(total, qx):\n", " return total + jnp.sum(dequantize(qx)), ()\n", "\n", "total, () = jax.lax.scan(sum_dequantized, 0., qxs, length=4)\n", "print(total)" ] }, { "cell_type": "markdown", "id": "18e3c7e6", "metadata": {}, "source": [ "Hi values also work as stacked outputs and as the loop carry. Here the\n", "carry is a quantized array, re-quantized after each accumulation step, and\n", "the second output stacks one fresh `QArray` per step into a `q8[4,2,3]`:" ] }, { "cell_type": "code", "execution_count": null, "id": "bbcb72fc", "metadata": {}, "outputs": [], "source": [ "def accum_quantized(qtotal, x):\n", " return quantize(dequantize(qtotal) + x), quantize(2 * x)\n", "\n", "qzero = quantize(jnp.zeros((2, 3), 'float32'))\n", "qtotal, qys = jax.lax.scan(accum_quantized, qzero, xs)\n", "print(jax.typeof(qtotal))\n", "print(jax.typeof(qys))" ] }, { "cell_type": "markdown", "id": "270bc322", "metadata": {}, "source": [ "## Sharding in types: explicit mode\n", "\n", "Finally, sharding. In JAX's explicit sharding mode (see [the parallelism\n", "guide](https://docs.jax.dev/en/latest/parallel.html)), shardings are part\n", "of array *types*: `jax.typeof` reports how a value is partitioned across\n", "the mesh, sharding propagation happens while tracing, and mismatches\n", "surface as type errors. The `sharding` field on `QArrayTy`, and the typing\n", "rules on our primitives that propagate it, are exactly what let hi types\n", "participate. What does it mean to partition a quantized array across\n", "devices? The components move together: if we shard the rows, each device\n", "holds its rows' quantized values *and* their scales. As usual there's a\n", "design choice to make, and we made it back in `lo_ty`: `qvalue` carries\n", "the type's sharding, and `scale` shards like it with the last axis\n", "dropped, so every row travels with its scale.\n", "\n", "Let's see it work. We make a mesh (this is where we use the 8 CPU devices\n", "requested in this document's first cell), shard some rows across it, and\n", "quantize — the shardings propagate through our typing rules into the\n", "result type, which `jax.typeof` displays with `@` markers:" ] }, { "cell_type": "code", "execution_count": null, "id": "57b5d4a0", "metadata": {}, "outputs": [], "source": [ "mesh = jax.make_mesh((4,), ('i',))\n", "jax.set_mesh(mesh)\n", "\n", "rows = jax.device_put(jnp.arange(24., dtype='float32').reshape(8, 3),\n", " jax.P('i'))\n", "\n", "qrows = quantize(rows)\n", "print(jax.typeof(qrows))\n", "print(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)\n", "print(jax.typeof(dequantize(qrows)))" ] }, { "cell_type": "markdown", "id": "066ed148", "metadata": {}, "source": [ "The same holds while tracing — hi types in jaxprs now display their\n", "shardings, computed by our `out_aval` rules:" ] }, { "cell_type": "code", "execution_count": null, "id": "61beaf4a", "metadata": {}, "outputs": [], "source": [ "jax.jit(lambda x: dequantize(quantize(x))).trace(rows).jaxpr" ] }, { "cell_type": "markdown", "id": "a25c4969", "metadata": {}, "source": [ "The quantized matmul propagates shardings too: its typing rule partitions\n", "output rows like `x`'s rows and output columns like `qw`'s columns\n", "(requiring the contracted axes to be unsharded, since neither operand can\n", "say how a sharded contraction should land):" ] }, { "cell_type": "code", "execution_count": null, "id": "760f5c19", "metadata": {}, "outputs": [], "source": [ "w2 = jnp.arange(12., dtype='float32').reshape(3, 4) / 12.\n", "\n", "print(jax.typeof(matmul_q(rows, quantize(w2)))) # rows sharded\n", "\n", "qw2 = quantize(jax.device_put(w2, jax.P(None, 'i')))\n", "print(jax.typeof(qw2)) # cols sharded\n", "print(jax.typeof(matmul_q(jnp.ones((2, 3), 'float32'), qw2)))" ] }, { "cell_type": "markdown", "id": "63409cd8", "metadata": {}, "source": [ "And the contraction check fires as a type error, at trace time:" ] }, { "cell_type": "code", "execution_count": null, "id": "faa4a663", "metadata": {}, "outputs": [], "source": [ "xk = jax.device_put(jnp.ones((2, 8), 'float32'), jax.P(None, 'i'))\n", "qk = quantize(jax.device_put(jnp.ones((8, 3), 'float32'), jax.P('i')))\n", "\n", "try:\n", " matmul_q(xk, qk)\n", "except TypeError as e:\n", " print('TypeError:', e)" ] }, { "cell_type": "markdown", "id": "93462e0d", "metadata": {}, "source": [ "Autodiff composes with all of this. Recall that `MatmulQ`'s backward rule\n", "passed `out_sharding` hints to its matmuls: that's because the cotangent\n", "for `qw` contracts over the row axis, which may be sharded (as it is\n", "here), and explicit mode refuses to guess how an all-sharded contraction\n", "should land. The right answer is the primal operand's sharding —\n", "cotangents live where their primals live:" ] }, { "cell_type": "code", "execution_count": null, "id": "a0b707c7", "metadata": {}, "outputs": [], "source": [ "def qloss(x, qw):\n", " return jnp.sum(matmul_q(x, qw) ** 2)\n", "\n", "grad_rows, grad_qw2 = jax.grad(qloss, argnums=(0, 1))(rows, quantize(w2))\n", "print(jax.typeof(grad_rows), jax.typeof(grad_qw2))" ] }, { "cell_type": "markdown", "id": "5acca0e9", "metadata": {}, "source": [ "One caution: JAX does not cross-check a hi primitive's declared output\n", "sharding against what its `expand` actually produces. The declared\n", "`out_aval` is what downstream code sees while tracing; `expand` determines\n", "what happens at runtime. Keeping them consistent is part of the typing\n", "rule's job.\n", "\n", "## `shard_map` and partition specs\n", "\n", "Explicit mode partitions values while keeping one global view of the\n", "program. Its complement is `shard_map`, which gives a per-device view. To\n", "cross that boundary, a quantized array needs one more piece: a partition\n", "spec type of our own, saying how `shard_map`'s `in_specs`/`out_specs`\n", "apply to it. We keep the same design as above: a quantized array is\n", "sharded along its leading axes only.\n", "\n", "For `shard_map`, we express this with an `HiPspec` subclass — the\n", "partition spec analogue of the `MappingSpec` above. Users pass instances\n", "of it as `in_specs`/`out_specs` entries, and its `to_lo` method says how\n", "it translates to one `jax.P` partition spec per lowered component (in\n", "`lo_ty` order):" ] }, { "cell_type": "code", "execution_count": null, "id": "4e4bd5e0", "metadata": {}, "outputs": [], "source": [ "from jax.experimental.hijax import HiPspec\n", "\n", "@dataclass(frozen=True)\n", "class QArrayP(HiPspec):\n", " spec: jax.P # partitioning of the leading axes; the last axis stays whole\n", "\n", " def to_lo(self):\n", " return (self.spec, self.spec) # qvalue and scale shard together" ] }, { "cell_type": "markdown", "id": "393fb202", "metadata": {}, "source": [ "(Since `scale` has one axis fewer than `qvalue`, handing both components\n", "the same partition spec says exactly that the leading axes shard together\n", "while the trailing axis of `qvalue` is untouched.)\n", "\n", "On the type, `shard` and `unshard` compute the per-device shard type from\n", "the global type and vice versa, delegating to the component types:" ] }, { "cell_type": "code", "execution_count": null, "id": "e254a848", "metadata": {}, "outputs": [], "source": [ "def qarray_shard(self, mesh, manual_axes, check_vma, spec):\n", " qvalue_ty, _ = self.lo_ty()\n", " qspec, _ = spec.to_lo()\n", " shard_ty = qvalue_ty.shard(mesh, manual_axes, check_vma, qspec)\n", " return QArrayTy(shard_ty.shape, shard_ty.sharding)\n", "\n", "def qarray_unshard(self, mesh, check_vma, spec):\n", " qvalue_ty, _ = self.lo_ty()\n", " qspec, _ = spec.to_lo()\n", " full_ty = qvalue_ty.unshard(mesh, check_vma, qspec)\n", " return QArrayTy(full_ty.shape, full_ty.sharding)\n", "\n", "QArrayTy.shard = qarray_shard\n", "QArrayTy.unshard = qarray_unshard" ] }, { "cell_type": "markdown", "id": "2ccb9aa6", "metadata": {}, "source": [ "Now quantized arrays can cross `shard_map` boundaries in either\n", "direction, continuing with the mesh and `rows` from the previous section.\n", "Producing a sharded quantized array:" ] }, { "cell_type": "code", "execution_count": null, "id": "ecd5c343", "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "@jax.shard_map(in_specs=jax.P('i'), out_specs=QArrayP(jax.P('i')))\n", "def quantize_shards(x):\n", " assert jax.typeof(x).shape == (2, 3) # each device sees two rows\n", " return quantize(x)\n", "\n", "qrows = quantize_shards(rows)\n", "print(jax.typeof(qrows))\n", "print(qrows.qvalue.sharding.spec, qrows.scale.sharding.spec)" ] }, { "cell_type": "markdown", "id": "4f3ad841", "metadata": {}, "source": [ "And consuming one, where each device sees a per-shard `QArray` of its own\n", "rows:" ] }, { "cell_type": "code", "execution_count": null, "id": "321f04ab", "metadata": {}, "outputs": [], "source": [ "@jax.jit\n", "@jax.shard_map(in_specs=QArrayP(jax.P('i')), out_specs=jax.P('i'))\n", "def dequantize_shards(qx):\n", " assert jax.typeof(qx).shape == (2, 3) # a per-device QArray shard\n", " return dequantize(qx)\n", "\n", "print(jnp.max(jnp.abs(dequantize_shards(qrows) - rows)))" ] }, { "cell_type": "markdown", "id": "9db30b9c", "metadata": {}, "source": [ "Because scales are per-row, quantizing shard-by-shard agrees exactly with\n", "quantizing globally — the same property that made batching pleasant:" ] }, { "cell_type": "code", "execution_count": null, "id": "767a85ef", "metadata": {}, "outputs": [], "source": [ "qrows_global = quantize(rows)\n", "assert (qrows.qvalue == qrows_global.qvalue).all()\n", "assert (qrows.scale == qrows_global.scale).all()" ] }, { "cell_type": "markdown", "id": "5eea947e", "metadata": {}, "source": [ "(For autodiff *through* a `shard_map`, there's a bit more to implement:\n", "`to_tangent_spec` and `to_ct_spec` on the spec type, and `nospec` on the\n", "hi type, which is used to shard autodiff residuals.)\n", "\n", "## Additional examples\n", "\n", "The recipe is always the same — a value class, a `HiType` with the\n", "`lo_ty`/`lower_val`/`raise_val` lowering triple, and primitives whose type\n", "signatures mention the new type — so further examples can be read quickly.\n", "\n", "### Rank-1 arrays\n", "\n", "A rank-1 array represents an `m × n` matrix as an outer product of two\n", "vectors, `col : f32[m]` and `row : f32[n]`. The point of the\n", "representation is to never materialize the `m × n` product: ops consume\n", "the factors directly. (A general low-rank type, with `f32[m, r]` and\n", "`f32[r, n]` factors, is more of the same.)" ] }, { "cell_type": "code", "execution_count": null, "id": "0786d2b9", "metadata": {}, "outputs": [], "source": [ "@dataclass(frozen=True)\n", "class Rank1:\n", " col: jax.Array # f32[m]\n", " row: jax.Array # f32[n]\n", "\n", "@dataclass(frozen=True)\n", "class Rank1Ty(HiType):\n", " shape: tuple[int, int] # (m, n), the dense shape represented\n", " sharding: NamedSharding # sharding of the dense shape; the factors' derive\n", "\n", " def lo_ty(self):\n", " (m, n), spec = self.shape, self.sharding.spec\n", " return [ShapedArray((m,), jnp.dtype('float32'),\n", " sharding=self.sharding.update(spec=jax.P(spec[0]))),\n", " ShapedArray((n,), jnp.dtype('float32'),\n", " sharding=self.sharding.update(spec=jax.P(spec[1])))]\n", " def lower_val(self, r1):\n", " return [r1.col, r1.row]\n", " def raise_val(self, col, row):\n", " return Rank1(col, row)\n", "\n", " # rank-1 matrices aren't closed under addition, so tangents are dense\n", " def to_tangent_aval(self):\n", " return ShapedArray(self.shape, jnp.dtype('float32'),\n", " sharding=self.sharding)\n", "\n", " def str_short(self, short_dtypes=False, mesh_axis_types=False):\n", " dims = [str(d) if p is None else f'{d}@{p}'\n", " for d, p in zip(self.shape, self.sharding.spec)]\n", " return f'r1[{\",\".join(dims)}]'\n", " __repr__ = str_short\n", "\n", "def typeof_rank1(r1):\n", " col_s, row_s = jax.typeof(r1.col).sharding, jax.typeof(r1.row).sharding\n", " sharding = col_s.update(spec=jax.P(col_s.spec[0], row_s.spec[0]))\n", " return Rank1Ty((r1.col.shape[0], r1.row.shape[0]), sharding)\n", "\n", "register_hitype(Rank1, typeof_rank1)" ] }, { "cell_type": "markdown", "id": "3f775110", "metadata": {}, "source": [ "The type records the dense shape it represents and prints as e.g.\n", "`r1[6,5]`. Note the tangent-type choice, which differs from `QArrayTy`'s\n", "for a different reason: the sum of two rank-1 matrices is in general\n", "rank 2, so rank-1 matrices aren't closed under addition and can't serve\n", "as their own tangent space. Perturbations leave the manifold, and\n", "tangents and cotangents are dense. The sharding story is the same as\n", "before, with the field consumed in `lo_ty`: the dense row axis is carried\n", "by `col` and the dense column axis by `row`.\n", "\n", "Unlike `QArray`, whose values were only ever created inside `quantize`,\n", "here construction from factors is itself a primitive — and there's an\n", "accessor primitive too, so that *rules* (which are ordinary traced code)\n", "can get at the factors without touching container attributes:" ] }, { "cell_type": "code", "execution_count": null, "id": "1b1b1ea8", "metadata": {}, "outputs": [], "source": [ "class Outer(HiPrim):\n", " def __init__(self, col_aval, row_aval):\n", " if not (len(col_aval.shape) == 1 and len(row_aval.shape) == 1):\n", " raise TypeError(f'bad outer factor types: {col_aval}, {row_aval}')\n", " sharding = col_aval.sharding.update(\n", " spec=jax.P(col_aval.sharding.spec[0], row_aval.sharding.spec[0]))\n", " self.in_avals = (col_aval, row_aval)\n", " self.out_aval = Rank1Ty((col_aval.shape[0], row_aval.shape[0]), sharding)\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, col, row):\n", " return Rank1(col, row)\n", "\n", " def vjp_fwd(self, nzs_in, col, row):\n", " return self(col, row), (col, row)\n", "\n", " def vjp_bwd_retval(self, res, g): # g is dense, f32[m, n]\n", " col, row = res\n", " return g @ row, col @ g\n", "\n", "class Factors(HiPrim):\n", " def __init__(self, r1_aval):\n", " self.in_avals = (r1_aval,)\n", " self.out_aval = tuple(r1_aval.lo_ty()) # the factor types are the lo types\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, r1):\n", " return (r1.col, r1.row)\n", "\n", "class MatmulR1(HiPrim):\n", " def __init__(self, x_aval, r1_aval):\n", " if not (isinstance(r1_aval, Rank1Ty) and len(x_aval.shape) == 2 and\n", " x_aval.shape[1] == r1_aval.shape[0]):\n", " raise TypeError(f'bad matmul_r1 operand types: {x_aval} @ {r1_aval}')\n", " self.in_avals = (x_aval, r1_aval)\n", " out_sharding = x_aval.sharding.update(\n", " spec=jax.P(x_aval.sharding.spec[0], r1_aval.sharding.spec[1]))\n", " self.out_aval = ShapedArray((x_aval.shape[0], r1_aval.shape[1]),\n", " jnp.dtype('float32'), sharding=out_sharding)\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, x, r1):\n", " return jnp.outer(x @ r1.col, r1.row) # never materialize col ⊗ row\n", "\n", " def vjp_fwd(self, nzs_in, x, r1):\n", " col, row = factors(r1) # rules are traced code: primitives, not attributes\n", " return self(x, r1), (x, col, row)\n", "\n", " def vjp_bwd_retval(self, res, g):\n", " x, col, row = res\n", " return jnp.outer(g @ row, col), x.T @ g\n", "\n", "def outer(col, row):\n", " return Outer(jax.typeof(col), jax.typeof(row))(col, row)\n", "\n", "def factors(r1):\n", " return Factors(jax.typeof(r1))(r1)\n", "\n", "def matmul_r1(x, r1):\n", " return MatmulR1(jax.typeof(x), jax.typeof(r1))(x, r1)" ] }, { "cell_type": "markdown", "id": "88068f9f", "metadata": {}, "source": [ "(Note `Factors` declares its output type as a *tuple* of types —\n", "`in_avals` entries and `out_aval` can be pytrees of types — and that it\n", "carries no autodiff rules at all, since we only apply it in forward\n", "passes; rules are only needed for the transformations you actually use.\n", "Note also `MatmulR1`'s backward rule computes the dense operand's\n", "cotangent as an outer product, exploiting the representation the same way\n", "`expand` does.)" ] }, { "cell_type": "code", "execution_count": null, "id": "c91c604d", "metadata": {}, "outputs": [], "source": [ "col = jnp.arange(6., dtype='float32') / 6.\n", "row = jnp.arange(5., dtype='float32') / 5.\n", "r1 = outer(col, row)\n", "print(jax.typeof(r1))\n", "\n", "acts = jnp.ones((3, 6), 'float32')\n", "print(jnp.max(jnp.abs(matmul_r1(acts, r1) - acts @ jnp.outer(col, row))))" ] }, { "cell_type": "markdown", "id": "bd4501ba", "metadata": {}, "source": [ "Jaxprs, `jit`, and autodiff all work as before. The gradient with respect\n", "to a rank-1 operand is dense, as the tangent type dictates, while\n", "gradients chained through the constructor land back on the factors:" ] }, { "cell_type": "code", "execution_count": null, "id": "36e52933", "metadata": {}, "outputs": [], "source": [ "jax.jit(matmul_r1).trace(acts, r1).jaxpr" ] }, { "cell_type": "code", "execution_count": null, "id": "cb7dae19", "metadata": {}, "outputs": [], "source": [ "def r1_loss(x, r1):\n", " return jnp.sum(matmul_r1(x, r1) ** 2)\n", "\n", "print(jax.typeof(jax.grad(r1_loss, argnums=1)(acts, r1)))\n", "\n", "def factor_loss(col, row):\n", " return jnp.sum(matmul_r1(acts, outer(col, row)) ** 2)\n", "\n", "g_col, g_row = jax.grad(factor_loss, argnums=(0, 1))(col, row)\n", "print(jax.typeof(g_col), jax.typeof(g_row))" ] }, { "cell_type": "markdown", "id": "f21b2baf", "metadata": {}, "source": [ "And the typing rules propagate shardings, just like `QArrayTy`'s:" ] }, { "cell_type": "code", "execution_count": null, "id": "7f5c1679", "metadata": {}, "outputs": [], "source": [ "print(jax.typeof(outer(jax.device_put(jnp.zeros(8, 'float32'), jax.P('i')),\n", " row)))" ] }, { "cell_type": "markdown", "id": "21eaf90f", "metadata": {}, "source": [ "We stopped here, but `vmap`, `scan`, and `shard_map` support would follow\n", "the same recipes as in the sections above: `dec_rank`/`inc_rank` and a\n", "mapping spec, `leading_axis_spec`, and `shard`/`unshard` with an `HiPspec`\n", "partition spec type.\n", "\n", "### Tuples\n", "\n", "Our last example is a generic container: a tuple whose elements are any\n", "JAX values. Where `QArrayTy` and `Rank1Ty` had a fixed component\n", "structure, `TupTy` is parameterized by its component *types* — and every\n", "method delegates to them:" ] }, { "cell_type": "code", "execution_count": null, "id": "3826882f", "metadata": {}, "outputs": [], "source": [ "import itertools\n", "\n", "@dataclass(frozen=True)\n", "class HiTup:\n", " elts: tuple\n", "\n", "@dataclass(frozen=True)\n", "class TupTy(HiType):\n", " tys: tuple # component types: array types, or other hi types\n", "\n", " # lowering delegates to the component types\n", " def lo_ty(self):\n", " return [lo for ty in self.tys for lo in ty.lo_ty()]\n", " def lower_val(self, tup):\n", " return [lo for ty, elt in zip(self.tys, tup.elts)\n", " for lo in ty.lower_val(elt)]\n", " def raise_val(self, *los):\n", " los = iter(los)\n", " return HiTup(tuple(ty.raise_val(*itertools.islice(los, len(ty.lo_ty())))\n", " for ty in self.tys))\n", "\n", " # so does the tangent type\n", " def to_tangent_aval(self):\n", " return TupTy(tuple(ty.to_tangent_aval() for ty in self.tys))\n", "\n", " def str_short(self, short_dtypes=False, mesh_axis_types=False):\n", " return ('Tup{' +\n", " ','.join(t.str_short(short_dtypes, mesh_axis_types)\n", " for t in self.tys) + '}')\n", " __repr__ = str_short\n", "\n", "register_hitype(HiTup, lambda t: TupTy(tuple(map(jax.typeof, t.elts))))" ] }, { "cell_type": "markdown", "id": "dd1c378a", "metadata": {}, "source": [ "Two things fall out of the delegation for free. First, elements can\n", "themselves be hi types — tuples of tuples, or a tuple holding a `QArray` —\n", "since `lo_ty` and friends just recurse. Second, there's no `sharding`\n", "field this time: the stored component types carry their own shardings, a\n", "third way of handling sharding-in-types alongside `QArrayTy`'s single\n", "field and `Rank1Ty`'s dense-shape spec.\n", "\n", "The constructor and accessor primitives are familiar from the rank-1\n", "example. One new ingredient: the element index is a *static parameter*,\n", "passed via `params` and available as `self.idx`. And since each element\n", "can be batched along its own axis (or not at all), the mapping spec\n", "carries one entry per component:" ] }, { "cell_type": "code", "execution_count": null, "id": "c7d56a81", "metadata": {}, "outputs": [], "source": [ "@dataclass(frozen=True)\n", "class TupSpec(MappingSpec):\n", " val: tuple # one axis entry per component\n", "\n", "def tup_dec_rank(self, size, spec):\n", " return TupTy(tuple(ty.dec_rank(size, s)\n", " for ty, s in zip(self.tys, spec.val)))\n", "\n", "def tup_inc_rank(self, size, spec):\n", " return TupTy(tuple(ty.inc_rank(size, s)\n", " for ty, s in zip(self.tys, spec.val)))\n", "\n", "TupTy.dec_rank = tup_dec_rank\n", "TupTy.inc_rank = tup_inc_rank\n", "\n", "class MakeTup(HiPrim):\n", " def __init__(self, elt_avals):\n", " self.in_avals = tuple(elt_avals)\n", " self.out_aval = TupTy(tuple(elt_avals))\n", " self.params = {}\n", " super().__init__()\n", "\n", " def expand(self, *elts):\n", " return HiTup(elts)\n", "\n", " def batch(self, axis_data, args, in_dims):\n", " return make_tup(*args), TupSpec(tuple(in_dims))\n", "\n", "class GetTupElt(HiPrim):\n", " def __init__(self, tup_aval, idx):\n", " self.in_avals = (tup_aval,)\n", " self.out_aval = tup_aval.tys[idx]\n", " self.params = dict(idx=idx)\n", " super().__init__()\n", "\n", " def expand(self, tup):\n", " return tup.elts[self.idx]\n", "\n", " def batch(self, axis_data, args, in_dims):\n", " tup, = args\n", " spec, = in_dims\n", " if spec is None:\n", " return get_tuple_element(tup, self.idx), None\n", " return get_tuple_element(tup, self.idx), spec.val[self.idx]\n", "\n", "def make_tup(*elts):\n", " return MakeTup(map(jax.typeof, elts))(*elts)\n", "\n", "def get_tuple_element(tup, idx):\n", " return GetTupElt(jax.typeof(tup), idx)(tup)" ] }, { "cell_type": "markdown", "id": "fd492d6a", "metadata": {}, "source": [ "(Note `GetTupElt`'s batch rule handles unbatched inputs — hi primitive\n", "batch rules are invoked even when no argument is batched, so `in_dims`\n", "entries can be `None`.)\n", "\n", "Tuples work, they nest, and they hold other hi types:" ] }, { "cell_type": "code", "execution_count": null, "id": "dab97eee", "metadata": {}, "outputs": [], "source": [ "tup = make_tup(jnp.arange(3.), 5.)\n", "print(jax.typeof(tup))\n", "print(get_tuple_element(tup, 1))\n", "\n", "nested = make_tup(make_tup(1., 2.), jnp.arange(2.))\n", "print(jax.typeof(nested))\n", "print(jax.jit(lambda t: get_tuple_element(get_tuple_element(t, 0), 1))(nested))\n", "\n", "print(jax.typeof(make_tup(qx, 3.))) # a quantized array element" ] }, { "cell_type": "markdown", "id": "cb802df4", "metadata": {}, "source": [ "The per-component mapping spec is the payoff under `vmap`: each element\n", "gets its own `in_axes`/`out_axes` entry, visible in the types. Here the\n", "first element is mapped on the way in, and only the second on the way\n", "out:" ] }, { "cell_type": "code", "execution_count": null, "id": "dade3db6", "metadata": {}, "outputs": [], "source": [ "def swap(t):\n", " a, b = get_tuple_element(t, 0), get_tuple_element(t, 1)\n", " return make_tup(b, a)\n", "\n", "out = jax.vmap(swap, in_axes=TupSpec((0, None)), out_axes=TupSpec((None, 0)),\n", " axis_size=3)(tup)\n", "print(jax.typeof(tup), '->', jax.typeof(out))" ] }, { "cell_type": "markdown", "id": "8f406d1c", "metadata": {}, "source": [ "And since the component types carry their own shardings, sharding-in-types\n", "needs nothing extra at all:" ] }, { "cell_type": "code", "execution_count": null, "id": "d14c26e8", "metadata": {}, "outputs": [], "source": [ "print(jax.typeof(make_tup(rows, jnp.float32(1.))))" ] }, { "cell_type": "markdown", "id": "5f672014", "metadata": {}, "source": [ "For autodiff rules on these primitives, and `scan` and `shard_map`\n", "support (via a per-component `HiPspec`), see the `TupTy` example in\n", "`tests/hijax_test.py`.\n", "\n", "## What we haven't covered\n", "\n", "On the primitive side, there are also hooks for customizing\n", "rematerialization and dead code elimination.\n", "\n", "As ever with hijax, `tests/hijax_test.py` is a good source of worked\n", "examples, and {ref}`hijax-custom-derivatives` covers the primitive-side\n", "API — including JVP rules, symbolic zeros, and custom linearization — in\n", "more depth." ] } ], "metadata": { "jupytext": { "cell_metadata_filter": "-all", "formats": "ipynb,md:myst,py:light", "main_language": "python", "notebook_metadata_filter": "nosearch" }, "nosearch": true }, "nbformat": 4, "nbformat_minor": 5 }