Article URL: https://docs.jax.dev/en/latest/hijax_types.html Comments URL: https://news.ycombinator.com/item?id=48884100 Points: 10 # Comments: 2

JAX’s built-in currency is the array: functions you transform take arrays in and produce arrays out, and every intermediate the tracing machinery sees has an array type like f32[3,4]. When you want to work with aggregate data, the usual tool is a pytree: you bundle arrays into containers, and JAX transparently flattens the bundle into its array leaves at every boundary. But sometimes transparency is exactly what you don’t want. Some data is best modeled as a new type, with its own identity: it should appear in jaxprs as a single value of a single type, not as a spray of array leaves; it has internal invariants, so users should only produce and consume it through a fixed set of operations, rather than by freely constructing or pattern-matching its components; its tangent type may differ from its primal structure, so that derivatives with respect to it aren’t just “the same pytree, but for tangents”; it can carry sharding information in the type, participating in JAX’s explicit sharding mode. Hijax types (or “hi types”) provide this. You subclass HiType to define the type, register a Python class as carrying values of that type, and write hijax primitives whose input and output types mention the new type. This document walks through the whole story with one running example: a quantized array type. We’ll assume some familiarity with hijax primitives; see Custom derivative rules with hijax primitives for an introduction to them. Like everything hijax, this is experimental: expect imports from jax.experimental.hijax, and expect the APIs to evolve.