JAX-style ML in the browser, with manual memory bookkeeping

What it does

jax-js is a from-scratch JavaScript machine learning framework that exposes a NumPy and JAX-compatible API for the browser. It compiles array operations into a custom intermediate representation, then synthesizes kernels for WebGPU, WebGL, and WebAssembly backends. The whole thing ships in an 80 KB gzip bundle with zero external dependencies, making GPU-accelerated ML about as portable as a browser tab.

The interesting bit

The standout design choice is a manual reference-counting memory model. Because JavaScript lacks deterministic destructors, the library forces you to track array ownership with .ref and .dispose(), and core combinators like jit and vmap assume your own functions obey the same contract. It is a rare framework that asks you to become the garbage collector.

Key highlights

• Synthesizes CPU kernels via Wasm and GPU kernels via WebGPU/WebGL from a shared compiler IR.
• Supports JAX-style composable transformations: jit, grad, vmap, and forward-mode autodiff.
• Ships at 80 KB gzipped per the project’s own comparison table, far smaller than TensorFlow.js or ONNX Runtime Web.
• Runs fully offline; no server-side Python or model hosting is required.
• Includes working demos for MNIST training, CLIP embeddings, voice cloning, and fluid simulation.

Caveats

• Memory management is manual: because JavaScript has no destructors, you must increment and decrement array references yourself, and the autodiff and JIT combinators assume your code follows these ownership rules.
• WebGPU support has platform holes: Node.js and Firefox for Android lack it entirely, and Safari support is gated to specific OS versions.
• SVD and BFloat16 are unsupported, and running ONNX models is only partially implemented.

Verdict

Worth a look if you need portable, lightweight in-browser ML with a familiar NumPy API and do not mind managing memory by hand. If you need Node.js GPU acceleration or expect Python-level ergonomics, this is not yet a drop-in replacement.