jax-js: JAX-Style Autodiff and Kernel Fusion That Runs in a Browser Tab

What jax-js Is For, and Who Should Care

The README describes jax-js as a machine learning framework for the browser that aims to bring JAX-style, high-performance CPU and GPU kernels to JavaScript. The problem it addresses is specific: numerical code that belongs on a GPU has historically meant Python plus CUDA, or a JavaScript runtime that only executes pre-exported graphs. jax-js takes the third path. You write array code in JavaScript, and the library compiles it.

The intended audience is narrow but real. If you are porting a JAX notebook to a web demo, or building an interactive visualization where the user adjusts parameters and expects the result immediately, the round trip to a Python backend is the bottleneck you are trying to remove. The examples listed in the README point the same way: an in-browser REPL, a fluid simulation, the heat method for geodesic distances, neural cellular automata, PCA, a Mandelbrot set renderer. These are programs where the compute is continuous and the latency budget is one frame, not one HTTP request.

The secondary audience is people who want autodiff without Python. The feature table claims gradients, Jacobian and Hessian, jvp() forward differentiation, jit() kernel fusion and vmap() auto-vectorization. TensorFlow.js is listed with gradients only, and onnxruntime-web with none of them. That gap is the core of the pitch.

The Compilation Path: Arrays In, Wasm and WebGPU Kernels Out

The README states that jax-js translates array operations into a compiler representation, then synthesizes kernels in WebAssembly and WebGPU. That two-stage design explains most of the library's behaviour. Array operations are not executed eagerly against a fixed set of hand-written kernels. They are captured, and the kernels are generated for the target backend.

This is why jit() and vmap() can exist at all. Kernel fusion requires a graph to fuse, and auto-vectorization requires a loop to rewrite. A library that dispatches each operation to a precompiled shader has neither. The README also lists graph capture as supported, alongside WebGPU, WebGL and Wasm backends, and mixed devices and mixed precision. Mixed devices is the unusual one: it implies you can hold tensors on different backends within one program, which matters when part of a computation is cheap on CPU and part is not.

The library is written from scratch with zero external dependencies. In practice that means the bundle does not pull in a numerical runtime, and the 80 KB gzip figure in the comparison table is plausible for a library that ships its own compiler rather than wrapping one. It also means numerical edge cases are the project's responsibility, not inherited from a mature C library. The table claims NaN/Inf numerics and Float64 support, which TensorFlow.js is listed as lacking. Float64 in a browser is not free: WebGPU has no native double precision, so a Float64 path either emulates it or falls back to the CPU backend. The material does not say which, and that is worth checking before you build on it.

Installing and Reaching the Library From a Page

The README gives two entry points. With a package manager:

npm i @jax-js/jax

The import is namespaced, and the README's quickstart uses an alias so the code reads like NumPy:

import { numpy as np } from "@jax-js/jax"; const x = np.array([1, 2, 3]); const y = x.mul(4); // [4, 8, 12]

For a blank HTML page with no bundler, the README shows a module script tag pulling from a CDN:

<script type="module"> import { numpy as np } from "https://esm.sh/@jax-js/jax"; </script>

There is no initialization call, no backend registration and no device selection shown in the quickstart. That is a deliberate contrast with runtimes that require you to await a session or configure an execution provider before the first tensor exists. The trade-off is that backend selection is implicit. The platform table lists CPU (Wasm), GPU (WebGPU) and GPU (WebGL) as separate columns, and the library presumably picks among them, but the supplied material does not document how to force one. If you need to pin a backend for reproducible output, verify that before shipping.

The release history is worth noting for upgrade planning: jax/v0.1.24 shipped 2026-08-31, following v0.1.23 on 2026-08-17 and v0.1.22 on 2026-08-13. Three patch releases inside a month, all still on 0.1.x. That is an active project with an unstable API surface.

Platform Coverage Is the Real Constraint

The platform table is the most useful page in the README, and it is also the one most likely to be skimmed. WebGPU is marked available on Chrome and Edge, on Chrome for Android, and on Firefox, Safari and iOS only from macOS 26 or iOS 26 onward. Firefox for Android shows no WebGPU support at all. Node.js shows CPU only, no GPU. Deno shows WebGPU as async.

Read that as a deployment map. If your users are on Chrome, you get the full GPU path today. If they are on Safari or Firefox on older macOS, they fall back to WebGL or to the Wasm CPU backend, and the performance characteristics change completely. If your code runs in a Node.js service, you get Wasm and nothing else, which removes most of the reason to choose this library over a Python stack.

The README notes that WebGPU has gained wide support in browsers as of late 2025, which is true as a trend but does not change the version floors in the table. The library also claims to be likely the most portable GPU ML framework, since it runs anywhere a browser can run. That claim is defensible given the WebGL fallback column, but portability here means the code runs, not that it runs at the same speed. A WebGL fallback and a WebGPU path are different execution environments, and the comparison table's single Speed row does not distinguish between them.

Where jax-js Is the Wrong Tool

The comparison table marks Run ONNX models as partial. That is the clearest signal about fit. If your task is to take a model someone else trained, export it, and run inference, onnxruntime-web is built for exactly that, with a full ONNX path and a 90 KB gzip runtime plus a 24 MB Wasm binary. jax-js asks you to express the model as array operations instead. For a hand-written network that is fine. For a graph with hundreds of fused nodes produced by an exporter, it is a rewrite.

Two format gaps compound this. BFloat16 is marked unsupported, and packed Uint8 is marked unsupported, while onnxruntime-web lists packed Uint8 as partial and Float16 as supported. Quantized or bfloat16 checkpoints are common in the model weights people actually download. The README does list safetensors reading as supported, which TensorFlow.js and onnxruntime-web are both listed as lacking, so weight loading from that format is a genuine advantage. But reading the file is not the same as running the dtype inside it.

The third limitation is structural. The library is at 0.1.x with releases landing every two weeks. The API reference lives on a separate site, and the README points to a Compatibility Table in FEATURES.md rather than inlining it. Any production dependency on this library means tracking that table across patch versions. There is no stated deprecation policy in the supplied material, and no LTS branch is mentioned.

How It Differs From TensorFlow.js and onnxruntime-web

The three libraries in the README's comparison table are not variations on one idea. They are three different answers to what a browser ML runtime is.

onnxruntime-web executes static ONNX graphs. The model is fixed at load time, the runtime is an interpreter for a format, and the 24 MB Wasm binary reflects the cost of supporting the full operator set. There is no autodiff, because there is nothing to differentiate: the graph is already built. If you want to train in the browser, this is the wrong architecture.

TensorFlow.js offers an eager array API and gradients, in a TensorFlow-shaped syntax. It is listed as supporting WebGPU as a preview and as lacking jit(), vmap(), jvp(), Jacobian and Hessian. The difference from jax-js is not speed, it is whether the runtime can see a whole computation and rewrite it. jax-js captures operations into a compiler representation and synthesizes kernels from that; TensorFlow.js is listed without kernel fusion, which means each op runs as its own dispatch.

jax-js sits between them. It has the eager array API of TensorFlow.js and the compilation story of onnxruntime-web, with autodiff added. The cost is that you write the model rather than importing it. The comparison table lists jax-js as fastest on speed, but that row has no methodology attached in the supplied material, and speed claims without a stated benchmark configuration should be treated as a starting hypothesis to test on your own hardware rather than a settled result.

Licence, Maintenance and What a Patch Upgrade Costs You

The repository is MIT licensed, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are retained. That is a permissive licence with no copyleft obligation and no source disclosure requirement for your own code. It does not cover the licences of model weights you load through safetensors, which are separate and often not permissive. Nothing here is legal advice; if you are shipping a product, read the licence text and the licence of any weights you bundle.

The maintenance picture from the supplied data: last push 2026-09-10, not archived, releases at roughly two-week intervals through August 2026, all on the 0.1.x line. A pre-1.0 version number with that release cadence means the API can change between minor versions, and there is no documented upgrade guide in the material provided. Budget for reading the Compatibility Table and the release notes on each bump, and pin an exact version in package.json rather than using a caret range if you ship to production.

The upgrade cost is not only API churn. Because the library generates kernels, a change to the compiler can alter numerical output or backend selection without any change to your code. Pin the version, and keep a small numerical regression test that compares a few known tensor results against stored values. The README's benchmark suite and in-browser REPL are the two places to reproduce a suspected regression before filing it.

Editorial conclusion

Adopt jax-js if you are writing numerical or differentiable code that must execute in the browser, and you already think in NumPy or JAX terms. Do not adopt it if your workload is inference over an exported ONNX graph, if you need BFloat16 or packed Uint8, or if Node.js is your only runtime, since the compatibility table marks WebGPU as unavailable there. Before committing, open the Compatibility Table in the repository and confirm that every op in your model is listed, then run your largest tensor shape through the benchmark suite on the slowest device you intend to support, because the platform table ties WebGPU on Firefox, Safari and iOS to macOS 26 or iOS 26 and later.

Official sources

1. ekzhang/jax-js on GitHub
2. License: MIT
3. Project website
4. README
5. Releases