FrankenJAX — Explore

01 · What it is

What JAX is, and why it matters

JAX is a Python library for accelerator-oriented array computation and program transformation, built for high-performance numerical computing and large-scale machine learning. That is its own README's description, and it understates the interesting part. JAX's real idea is that the powerful operations are composable function transformations: grad takes the derivative of a function, jit compiles it to run fast on GPUs and TPUs via XLA, vmap vectorizes it over batches. You can stack them in any order. Jit of grad of vmap. A few lines of plain NumPy-like code become a compiled, differentiated, parallel program.

It matters because it sits underneath a large share of modern ML research (it is the engine room of Google's ML ecosystem), and because its semantics are defined by its implementation. There is no specification of what grad(vmap(f)) means outside the JAX codebase. The mathematics and the machinery are one entangled thing, written in Python, compiled through XLA, running on accelerators.

02 · Why build it

Why rewrite it in Rust

The maintainer — Jeffrey Emanuel, a solo developer building a whole suite of Rust reimplementations — argues that the mathematical core deserves to exist outside the machinery. Today there is no standalone, portable, verifiable implementation of JAX's transform semantics: no way to embed JAX-style automatic differentiation in a Rust application without dragging in Python and XLA, no reference implementation a verification engineer can audit, no canonical artifact that says what a transform composition means.

His distinctive addition is the Trace Transform Ledger. Every transform composition emits a proof artifact: the root program, the transform stack, the evidence, with a signature, so a pipeline can show its work. The audience is compiler researchers, verification engineers, and educators, not ML training. The README says plainly that this is not a replacement for JAX in production.

The counter-argument is that JAX owns this lane by definition, because the lane is JAX's own semantics. Against the incumbent, a reimplementation of the meaning of composition competes with the thing it references. That is a fight it cannot win on utility, and it does not try to win. The compilation half is explicitly out of scope: there is no XLA lowering, no GPU, no distributed execution. What exists is an interpreter. And the incumbent has its own portability answer: jax.export serializes lowered programs to StableHLO for execution without Python. Wherever the need is "run this traced program elsewhere," the export format already covers it. The remaining case is narrow: someone who needs JAX semantics without JAX the system (embedded Rust, air-gapped audit, verified pipelines), and who trusts conformance attested only by the author.

03 · What was built

What the project actually built

For one person, the build is substantial: 459,703 lines of Rust across 17 crates. A canonical IR. Tracing from Rust closures. An order-sensitive transform stack with dedicated property tests for each composition ordering. Tape-based reverse-mode and forward-mode automatic differentiation dispatched over 162 primitives. An e-graph optimizer. A differential conformance harness checked against 861 fixture cases captured from real JAX. The unsafe posture is the tightest our program has assessed: sixteen crates forbid unsafe code outright, and the single production unsafe block is the foreign-function call boundary itself, documented with a written safety argument.

The honest center of the project is also its strangest artifact. In this codebase, jit is a verified no-op. There is no compiler, so the dispatcher strips leading jit as a pass-through. A property test proves the identity to fourteen decimal places across random inputs. The project tests the transparency instead of hiding it. That is the shape of the actual scope: it reimplements the meaning of transform composition, not compilation.

Evidence · CI status

What the project’s own CI said at the pin

the pin = the commit the assessment was pinned to.

There is no public CI (continuous integration: automated checks that run on every change) at all: zero GitHub Actions workflows.

The uncomfortable findings

Uncomfortable finding

So every test count, gate result, and benchmark number is attested only by checked-in JSON produced on the maintainer's local fleet. The evidence machinery is elaborate: ledgers, forensic logs, gate scripts, proof matrices. It is all self-attested. The README contradicts itself and the code on the most basic facts: 118 versus 162 primitives, 15 versus 17 crates, line counts understated by roughly three times. The "157 of 157 hand-derived differentiation rules" claim rests partly on placeholder zero-gradient rules: fifty primitives sit behind zero rules, and several, like the matrix factorizations, are actually differentiable, so zero is a stand-in rather than a derivation. The GPU backend is a planning document, not code. The toolchain file the README points to is zero bytes, so no one can know which compiler the numbers were produced with. source source And the license is MIT plus a rider naming OpenAI and Anthropic, with their affiliates and agents, as forbidden parties, barring them even from analyzing the code, while 52 of the 100 most recent commits were co-authored by Anthropic's models, plus one by Grok.

04 · Where it stands

Where it stands

In our program's terms: TRL 4, NODUS ring Explore. In plain English: the differentiation core is implemented and the conformance discipline is real, but nothing has been observed executing outside the author's environment, the compilation half of the problem is explicitly out of scope, and the license rider means the program cannot benchmark, analyze, or evaluate the software while it stands. source Explore with a ceiling, and the ceiling is legal, not technical.

Technology readiness

13579

TRL 4 on a 1-to-9 scale.

NODUS ring

MonitorExplorePilotInvest

The program’s adoption ring for this project.

Verdict key

NODUS ring: Explore: The default for substantive-but-unproven work.
TRL: Technology Readiness Level, a 1-to-9 scale of how proven a technology is. This brief rates the project at TRL 4.
CI: Continuous integration: the project’s automated checks, run on every change.
The pin: The exact commit the assessment was pinned to. “At the pin” means as of that commit.

Evidence tiers

Tier 1 · Verified — Confirmed by direct inspection of a fresh clone, an API response, or a live page read by the analyst.
Tier 2 · CI-observed — Observed executing on live CI pages. Attests the suite runs, not that it is green, unless pass/fail is legible.
Tier 3 · Maintainer claim — Asserted in README or docs by the maintainer; not independently executed or reproduced.
Tier 4 · External — Independent sources: APIs, papers, press, third-party benchmarks. Absence of coverage is reported as a finding.
Tier 5 · Inference — The analyst’s judgment. Always labeled; never presented as fact.

What this verdict rests on. Tier labels are the assessment packet’s own annotations.

Tier 1 · Verified: Counts are static analysis from the clone; the AD core is implemented, not stubbed; zero GitHub Actions workflows — there is no CI to check.
Tier 3 · Maintainer claim: Nothing was compiled or executed; every test count and gate result is maintainer-asserted.
Tier 4 · External: Zero independent coverage found.
Tier 5 · Inference: The Explore ring.

Should you use it?

No.

There is no release, no compiler, and no independent witness that the tests pass.

Should you learn from it?

Yes.

The Trace Transform Ledger as an artifact class, the negative-evidence perf ledger as a self-audit template, and the differential-oracle discipline at scale are all exportable.

What would change the verdict, in order

the rider narrowed or removed (the necessary condition: nothing else matters without it), one public CI workflow running the test suite on a pinned toolchain, a release artifact, the fifty placeholder gradient rules audited into real derivations or typed errors

05 · What it teaches

What this teaches about building with agents

The transferable lesson is about what makes a proof artifact trustworthy: it must be able to reject.

The Trace Transform Ledger is the project's most novel idea: a signed record that a transform composition meant what it claimed. What separates it from decoration is the semantic proof matrix: twelve cases, six accepted and six correctly rejected. A ledger that can only say yes is a press release. One that demonstrably says no, on invalid compositions, is evidence. That is a design rule worth stealing for any agent-built system that needs to show its work. Build the rejection cases first, and publish the matrix.

The shadow lesson is the one the project has not yet learned. It built an elaborate self-attestation apparatus — ledgers, forensic logs, gate scripts, a 36,000-line evidence-regeneration commit — and attests everything to itself. Pride in the machinery confused the courtroom with the verdict. No independent party has ever been observed watching a test pass. The machinery is really good. It becomes trustworthy the day a second witness runs it. Until then, it is the author witnessing himself, very carefully.

FrankenJAX Rewriting JAX's Mathematical Core in Rust