jax-js: JAX in pure JavaScript

Website | API Reference | Compatibility Table | Discord

jax-js is a machine learning framework for the browser. It aims to bring JAX-style, high-performance CPU and GPU kernels to JavaScript, so you can run numerical applications on the web.

npm i @jax-js/jax

Under the hood, it translates array operations into a compiler representation, then synthesizes kernels in WebAssembly and WebGPU.

The library is written from scratch, with zero external dependencies. It maintains close API compatibility with NumPy/JAX. Since everything runs client-side, jax-js is likely the most portable GPU ML framework, since it runs anywhere a browser can run.

Quickstart

import { numpy as np } from "@jax-js/jax";

// Array operations, compatible with JAX/NumPy.
const x = np.array([1, 2, 3]);
const y = x.mul(4); // [4, 8, 12]

Web usage (CDN)

In vanilla JavaScript (without a bundler), just import from a module script tag. This is the easiest way to get started on a blank HTML page.

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

Platforms

This table refers to latest versions of each browser. WebGPU has gained wide support in browsers as of late 2025.

Platform	CPU (Wasm)	GPU (WebGPU)	GPU (WebGL)
Chrome / Edge	✅	✅	✅
Firefox	✅	✅ - macOS 26+	✅
Safari	✅	✅ - macOS 26+	✅
iOS	✅	✅ - iOS 26+	✅
Chrome for Android	✅	✅	✅
Firefox for Android	✅	❌	✅
Node.js	✅	❌	❌
Deno	✅	✅ - async	❌

Examples

Community usage:

• g9-jaxjs: Automatically interactive graphics with forward-mode AD
• autoresearch-webgpu: autoresesarch, in the browser
• tanh.xyz: Interactive ML visualizations
• handwritten: Handwriting synthesis with LSTM-MDN
• jax-js-bayes: Declarative Bayesian modeling library

Demos on the jax-js website:

• Chat with Gemma 3
• Training neural networks on MNIST
• Voice cloning: Kyutai Pocket TTS
• Speech recognition: OpenAI Whisper
• CLIP embeddings for books in-browser
• Object detection: D-FINE (ONNX)
• Object detection: DETR ResNet-50 (ONNX)
• Fluid simulation (Navier-Stokes)
• Neural cellular automata
• Principal component analysis
• Mandelbrot set
• In-browser REPL
• Benchmark suite

Feature comparison

Here's a quick, high-level comparison with other popular web ML runtimes:

Feature	jax-js	TensorFlow.js	onnxruntime-web
Overview
API style	JAX/NumPy	TensorFlow-like	Static ONNX graphs
Latest release	2026	⚠️ 2024	2026
Speed	Fastest	Fast	Fastest
Bundle size (gzip)	80 KB	269 KB	90 KB + 24 MB Wasm
Autodiff & JIT
Gradients	✅	✅	❌
Jacobian and Hessian	✅	❌	❌
jvp() forward differentiation	✅	❌	❌
jit() kernel fusion	✅	❌	❌
vmap() auto-vectorization	✅	❌	❌
Graph capture	✅	❌	✅
Backends & Data
WebGPU backend	✅	🟡 Preview	✅
WebGL backend	✅	✅	✅
Wasm (CPU) backend	✅	✅	✅
Eager array API	✅	✅	❌
Run ONNX models	🟡 Partial	❌	✅
Read safetensors	✅	❌	❌
Float64	✅	❌	❌
Float32	✅	✅	✅
Float16	✅	❌	✅
BFloat16	❌	❌	❌
Packed Uint8	❌	❌	🟡 Partial
Mixed precision	✅	❌	✅
Mixed devices	✅	❌	❌
Ops & Numerics
Arithmetic functions	✅	✅	✅
Matrix multiplication	✅	✅	✅
General einsum	✅	🟡 Partial	🟡 Partial
Sorting	✅	❌	❌
Activation functions	✅	✅	✅
NaN/Inf numerics	✅	✅	✅
Basic convolutions	✅	✅	✅
n-d convolutions	✅	❌	✅
Strided/dilated convolution	✅	✅	✅
Cholesky, Lstsq	✅	❌	❌
LU, Solve, Determinant	✅	❌	❌
Eigenvalues, SVD	✅	❌	❌
FFT	✅	🟡 Quadratic	🟡 Only on CPU
Basic RNG (Uniform, Normal)	✅	✅	✅
Advanced RNG	✅	❌	❌

Tutorial

Programming in jax-js looks very similar to JAX, just in JavaScript.

Arrays

Create an array with np.array():

import { numpy as np } from "@jax-js/jax";

const ar = np.array([1, 2, 3]);

By default, this is a float32 array, but you can specify a different dtype:

const ar = np.array([1, 2, 3], { dtype: np.int32 });

For more efficient construction, create an array from a JS TypedArray buffer:

const buf = new Float32Array([10, 20, 30, 100, 200, 300]);
const ar = np.array(buf).reshape([2, 3]);

Once you're done with it, you can unwrap a jax.Array back into JavaScript. This will also apply any pending operations or lazy updates:

// 1) Returns a possibly nested JavaScript array.
ar.js();
await ar.jsAsync(); // Faster, non-blocking

// 2) Returns a flat TypedArray data buffer.
ar.dataSync();
await ar.data(); // Fastest, non-blocking

Arrays can have mathematical operations applied to them. For example:

import { numpy as np, scipySpecial as special } from "@jax-js/jax";

const x = np.arange(100).astype(np.float32); // array of integers [0..99]

const y1 = x.ref.add(x.ref); // x + x
const y2 = np.sin(x.ref); // sin(x)
const y3 = np.tanh(x.ref).mul(5); // 5 * tanh(x)
const y4 = special.erfc(x.ref); // erfc(x)

Notice that in the above code, we used x.ref. This is because of the memory model, jax-js uses reference-counted ownership to track when the memory of an Array can be freed. More on this below.

Reference counting

Big Arrays take up a lot of memory. Python ML libraries override the __del__() method to free memory, but JavaScript has no such API for running object destructors (cf.). This means that you have to track references manually. jax-js tries to make this as ergonomic as possible, so you don't accidentally leak memory in a loop.

Every jax.Array has a reference count. This satisfies the following rules:

• Whenever you create an Array, its reference count starts at 1.
• When an Array's reference count reaches 0, it is freed and can no longer be used.
• Given an Array a:
• Accessing a.ref returns a and changes its reference count by +1.
• Passing a into any function as argument changes its reference count by -1.
• Calling a.dispose() also changes its reference count by -1.

What this means is that all functions in jax-js must take ownership of their arguments as references. Whenever you would like to pass an Array as argument, you can pass it directly to dispose of it, or use .ref if you'd like to use it again later.

You must follow these rules on your own functions as well! All combinators like jvp, grad, jit assume that you are following these conventions on how arguments are passed, and they will respect them as well.

// Bad: Uses `x` twice, decrementing its reference count twice.
function foo_bad(x: np.Array, y: np.Array) {
  return x.add(x.mul(y));
}

// Good: The first usage of `x` is `x.ref`, adding +1 to refcount.
function foo_good(x: np.Array, y: np.Array) {
  return x.ref.add(x.mul(y));
}

Here's another example:

// Bad: Doesn't consume `x` in the `if`-branch.
function bar_bad(x: np.Array, skip: boolean) {
  if (skip) return np.zeros(x.shape);
  return x;
}

// Good: Consumes `x` the one time in each branch.
function bar_good(x: np.Array, skip: boolean) {
  if (skip) {
    const ret = np.zeros(x.shape);
    x.dispose();
    return ret;
  }
  return x;
}

You can assume that every function in jax-js takes ownership properly, except with a couple of very rare exceptions that are documented.

grad(), vmap() and jit()

JAX's signature composable transformations are also supported in jax-js. Here is a simple example of using grad and vmap to compute the derivaive of a function:

import { numpy as np, grad, vmap } from "@jax-js/jax";

const x = np.linspace(-10, 10, 1000);

const y1 = vmap(grad(np.sin))(x.ref); // d/dx sin(x) = cos(x)
const y2 = np.cos(x);

np.allclose(y1, y2); // => true

The jit function is especially useful when doing long sequences of primitives on GPU, since it fuses operations together into a single kernel dispatch. This improves memory bandwidth usage on hardware accelerators, which is the bottleneck on GPU rather than raw FLOPs. For instance:

export const hypot = jit(function hypot(x1: np.Array, x2: np.Array) {
  return np.sqrt(np.square(x1).add(np.square(x2)));
});

Without JIT, the hypot() function would require four kernel dispatches: two multiplies, one add, and one sqrt. JIT fuses these together into a single kernel that does it all at once.

All functional transformations can take typed JsTree of inputs and outputs. These are similar to JAX's pytrees, and it's basically just a structure of nested JavaScript objects and arrays. For instance:

import { grad, numpy as np } from "@jax-js/jax";

type Params = {
  foo: np.Array;
  bar: np.Array[];
};

function getSums(p: Params) {
  const fooSum = p.foo.sum();
  const barSum = p.bar.map((x) => x.sum()).reduce(np.add);
  return fooSum.add(barSum);
}

grad(getSums)({
  foo: np.array([1, 2, 3]),
  bar: [np.array([10]), np.array([11, 12])],
});
// => { foo: [1, 1, 1], bar: [[1], [1, 1]] }

Note that you need to use type alias syntax rather than interface to define fine-grained JsTree types.

Devices

Similar to JAX, jax-js has a concept of "devices" which are a backend that stores Arrays in memory and determines how to execute compiled operations on them.

There are currently 4 devices in jax-js:

• cpu: Slow, interpreted JS, only meant for debugging.
• wasm: WebAssembly, multi-threaded when SharedArrayBuffer is available.
• webgpu: WebGPU, available on supported browsers (Chrome, Firefox, Safari, iOS).
• webgl: WebGL2, via fragment shaders. This is an older graphics API that runs on almost all browsers, but it is much slower than WebGPU. It's offered on a best-effort basis and not as well-supported.

We recommend webgpu for best performance, especially when running neural networks. T