July 12, 2026
Puns, shade, and one very nerdy showdown
Defining new Jax types with hijax
JAX gets a flashy new trick, and the comments instantly split between hype, confusion, and sass
TLDR: JAX’s new experimental hijax feature lets developers define custom data types so complex values can behave like one object instead of a messy bundle. Commenters loved the name, mocked the docs, and immediately questioned whether this solves a real problem or just competes with tools people already trust.
A new experimental add-on called hijax is giving JAX users a way to create their own custom data types instead of forcing everything to look like a pile of plain arrays. In normal-human terms: if you have something special, like a compressed or quantized number format, this lets it act like one real thing instead of two awkwardly taped-together parts. The pitch is all about stronger rules, cleaner internals, and making advanced math features behave more sensibly.
But let’s be honest: the community reaction is where the fun starts. One of the first replies was pure drive-by chaos: "the usual tool is a pytree: 404" from chrisjj, which instantly gave the whole thread a “docs error as performance art” vibe. Then came the branding applause, with 0xnyn basically declaring the name the real star of the show: "Hijacking JAX with hijax". If the feature felt intimidating, the pun at least landed.
And then the real mini-drama arrived: why use this at all if other tools already exist? peterdsharpe brought the skeptical energy, asking why anyone should pick this over Equinox’s Module plus jaxtyping, which they called the cleaner, more established option. That’s the fault line in the reactions: some see hijax as a powerful new way to keep complex data safer and more meaningful, while others are already side-eyeing it as another experimental layer in a very crowded toolbox. In classic dev-thread fashion, the launch became part feature reveal, part naming contest, part “do we really need this?” showdown.
Key Points
- •The article presents experimental hijax types as a way to define custom first-class types in JAX that appear as single values in jaxprs.
- •It argues that pytrees are insufficient when data must preserve invariants, expose a distinct type identity, or use custom tangent and batching semantics.
- •To define a hi type, developers subclass `HiType`, implement lowering and raising methods, and register a Python value class with `register_hitype`.
- •Custom hi-type values are produced and consumed through hijax primitives, including `VJPHiPrimitive` subclasses with typing and differentiation rules.
- •The running example is a quantized array `QArray` composed of an `int8` value tensor and a floating-point per-row scale, used to show why a pytree representation can be inadequate.