Skip to content
jax-js Brings JAX-Style ML Kernels To The Browser

jax-js Brings JAX-Style ML Kernels To The Browser

Eric Zhang’s jax-js is a pure JavaScript ML framework that compiles NumPy/JAX-style array code into WebAssembly and WebGPU kernels for client-side experiments.

Details

Published

What it is

jax-js is Eric Zhang’s attempt to make the browser feel a bit more like a tiny JAX lab. You write NumPy/JAX-style array code in JavaScript, then the library lowers those operations into its own compiler representation and emits kernels for WebAssembly and WebGPU. That is a pretty serious sentence for something you can import from a script tag.

The interesting part is not just that it runs ML-ish workloads client-side. It is the shape of the bet: numerical computing that travels with the page, works without a Python install, and can reach the GPU when the browser allows it. The examples are nicely varied too, from MNIST and matrix benchmarks to Gemma, Whisper, CLIP-style embeddings, fluid simulation, and neural cellular automata. If you like the idea of web pages that can compute instead of only display, this one is worth opening with the docs in the next tab.

How it’s made