Code Review

Hijax

jax.experimental.hijax is JAX's own extension API for custom types. It is complementary to quax rather than an alternative — the decision guide is in skills/quax/SKILL.md under "Quax, hijax, or both", and quax.examples.hijax is the worked combination. At review time:

• Ask whether it is needed at all. A hijax primitive costs a typing rule, an expand, and a rule per transform, for one operation. If the change could be a plain @quax.register rule instead, it should be. Hijax earns its cost only when the type must do something quax cannot: a cotangent type that differs from the primal's, invariants checked in the jaxpr, custom batching, or sharding in the type.
• ArrayValue.aval() must still return a ShapedArray. Returning a HiType makes every jnp function reject the tracer, because jax.Array's isinstance check reads the aval. Build the ShapedArray from the leaf's hijax type instead. tests/unit/test_hijax_interop.py pins this.
• A hijax value must not be made a pytree, and must not be constructed or attribute-accessed under a trace. Both are legal only inside expand and the type's own methods. A rule that reads value.array in traced code is a bug even though it works eagerly.
• Cotangent units/metadata must match to_ct_aval. JAX type-checks a bwd rule's output against it, so a rule that returns the primal's type where the cotangent type differs will fail at trace time rather than silently — but only if a test actually differentiates it.
• The API renames. Names must be resolved by probing, not pinned: VJPHiPrimitive becomes HiPrim after JAX 0.11.1 and MappingSpec only became public in 0.11.0. src/quax/experimental/hijax.py does this once, in _resolve; a change that spells either name anywhere else has reintroduced the coupling.
• Expect the JAX_CHECK_TRACER_LEAKS fixture, and check it stays narrow. Consuming a hi value under jit reports a false leak on JAX 0.10.2 and 0.11.0 only — a JAX regression, fixed in 0.11.1. tests/usage/test_hijax.py disables the check for those versions and is a no-op elsewhere. A PR that widens it to every version, weakens the check globally, or drops the version bound is a regression: the bound is what makes a recurrence visible.

The hot path

_trace.py, _dispatch.py, _module.py, and Value construction run once per primitive per call. Cost added here is multiplied by the size of the user's program, and this repo has spent four PRs (#168, #173, #174, #187) recovering it.

• Per-call object construction, isinstance against an ABC, or equinox per-instance validation in these files is a regression until benchmarked otherwise.
• Caches must not pin what they watch. #187 was a weakref-keyed cache holding a strong reference to the jaxpr it was tracking, so nothing was ever collected.
• CodSpeed benchmarks live in tests/benchmark/. A change plausibly affecting dispatch cost should say what happened to them.
• plum's resolution cache is disabled entirely by a single non-faithful type, which is why _compat.py sets jax.Array.__faithful__. Changes near dispatch registration can switch that cache off without any visible symptom but a slow benchmark.

Silent failure

Rule 3 of dispatch already degrades silently by design. Anything that makes an error path quieter compounds it and should be treated as a defect rather than a cleanup:

• assert False for an unreachable branch — assertions vanish under -O, and this was replaced with a real TypeError in #171.
• A bare except, or a fallback that swallows the reason it was reached.
• Broadening a rule's annotations so that a previously-unmatched call now materialises instead of raising.

Tests

The harness in tests/unit/test_lax/ and tests/unit/test_numpy/ runs each entry through both plain JAX and quax.quaxify, then asserts the results are equal. Value correctness for a covered primitive is therefore already checked. What it cannot check, and what review must:

• expect_myarray is the real assertion. It says whether the custom type survives the operation, per output. A True that should be False (or a tuple with the wrong arity) is a coverage hole the suite will happily pass.
• mark_todo is pytest.mark.skip. A PR that implements a primitive must remove the mark on its entry, or the new rule ships untested.
• Both files. test_myarray.py and test_jax_array.py are separate; a rule affecting the plain-array path needs an entry in each.
• Metadata semantics are invisible to MyArray, which carries no units or axis names. Bugs like #180 live in tests/usage/test_<name>.py — a rule on an example type needs coverage there, not only in the unit tests.
• The suite runs with JAX_CHECK_TRACER_LEAKS=1 and jaxtyping/beartype enabled.

Repo conventions

• uv run for everything — uv run pytest, uv run prek run --all-files. Never bare python or pytest.
• Commits use gitmoji plus conventional commits: 🐛 fix(trace): ..., ✅ test: ..., ⚡️ perf: ..., build(deps): ....
• Pre-commit runs pyright on src/ only, so type errors in tests/ reach CI unreviewed.
• Doctests are not collected. Examples in README.md, docs/, and src/ docstrings never run — a PR changing one has not tested it, and the diff needs reading on that basis. The exceptions are README.md and skills/quax/SKILL.md, whose python blocks are executed as individual tests by Sybil, wired up in conftest.py.

Further reading

• skills/quax/SKILL.md — the dispatch resolution ladder, writing rules, boundary behaviour, and a troubleshooting table.
• AGENTS.md — module-by-module architecture and commands.
• src/quax/examples/zero/_core.py — minimal reference type; lora/_core.py — advanced rules.