# jax-js: JAX-Style Autodiff and JIT in Pure JavaScript

> jax-js brings NumPy/JAX-compatible arrays, jit(), vmap() and gradients to the browser with WebAssembly and WebGPU backends. It is a good fit for client-side numerical work and a poor fit for server-side training.

**ekzhang/jax-js** — JAX in JavaScript – ML library for the web, running on WebGPU & Wasm

- Repository: https://github.com/ekzhang/jax-js
- Website: https://jax-js.com
- Stars: 936 · Forks: 61
- Language: TypeScript
- License: MIT
- Published: 2026-09-10 · Updated: 2026-09-10 · Language: en
- Canonical page: https://hysenlabs.com/projects/ekzhang-jax-js

## The problem jax-js solves: numerical code that never leaves the browser

Most machine learning runtimes assume a server. You ship tensors to a process, it holds the weights, and the browser is a thin client. jax-js inverts that. The README describes it as "a machine learning framework for the browser" that brings "JAX-style, high-performance CPU and GPU kernels to JavaScript", and it states that because everything runs client-side it is "likely the most portable GPU ML framework, since it runs anywhere a browser can run". The target reader is a JavaScript or TypeScript developer who wants JAX semantics without a Python service: someone building an interactive visualization, a browser demo, or a tool where shipping user data to a backend is undesirable. The API surface is deliberately familiar. The quickstart imports `numpy as np` from `@jax-js/jax`, and the README says the library "maintains close API compatibility with NumPy/JAX". If you already write `x.mul(4)` in JAX, the same expression works here, which lowers the cost of porting an existing numerical routine rather than rewriting it against a new tensor abstraction.

## How jax-js compiles array operations into Wasm and WebGPU kernels

The mechanism is a two-stage pipeline. The README states that under the hood jax-js "translates array operations into a compiler representation, then synthesizes kernels in WebAssembly and WebGPU". So array operations are not interpreted one by one against a runtime dispatch table. They are captured into an intermediate representation, and kernels are generated from it for the backend in use. That representation is what makes the JAX-style transformations possible. The feature table lists `jit()` kernel fusion, `vmap()` auto-vectorization, `jvp()` forward differentiation, Jacobian and Hessian computation, and graph capture as supported, while TensorFlow.js is marked as not supporting any of those. The project is written from scratch with, per the README, "zero external dependencies", which matters for bundle size: the table gives 80 KB gzip for jax-js against 269 KB for TensorFlow.js and 90 KB plus a 24 MB Wasm payload for onnxruntime-web. Backends are selected per platform, and the table shows a distinction that is easy to miss: on Node.js only the Wasm CPU backend is available, WebGPU is absent, and Deno supports WebGPU only asynchronously.

## Installing jax-js and running your first array

The README gives a single install command. Run it in a project that already has a package manager and a bundler, or skip to the CDN form below if you are working on a blank HTML page. After install, the package resolves through the exports map in package.json, which points `import` at `./dist/index.js` and `require` at `./dist/index.cjs`, with types at `dist/index.d.ts`.

```bash
npm i @jax-js/jax
```

The first real use is the quickstart from the README: import the NumPy namespace and do arithmetic on an array. The comment in the README shows the expected result, `[4, 8, 12]`, so if you log `y` you should see those three values.

```js
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]
```

If you do not want a build step, the README documents a module script tag that pulls the package from esm.sh. This is the path the README calls "the easiest way to get started on a blank HTML page", and it needs no npm install at all.

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

There is also an in-browser REPL linked from the README at jax-js.com/repl, which is the fastest way to check whether a specific operator behaves the way you expect before wiring it into an application.

## Where jax-js is the wrong tool: Node.js, BFloat16 and partial ONNX

The platforms table is the first constraint to read carefully. Node.js gets CPU via Wasm and nothing else: the WebGPU column is marked absent, and so is WebGL. Deno is marked as supporting WebGPU only asynchronously. Firefox for Android has no WebGPU entry. WebGPU on Firefox, Safari and iOS is listed as requiring macOS 26+ or iOS 26+ respectively. If your deployment target is a server process or an older mobile browser, the GPU path is not there and you are running on Wasm. The datatype table is the second constraint. Float64, Float32, Float16 and mixed precision are listed as supported. BFloat16 is not, and neither is packed Uint8. onnxruntime-web is marked as partially supporting packed Uint8, so if your model weights are quantized that way, this is not the runtime for you. The third constraint is model loading: the table marks "Run ONNX models" as partial for jax-js, while onnxruntime-web supports it fully. jax-js reads safetensors, which onnxruntime-web does not, so the two are complementary rather than interchangeable. Finally, the speed row in the comparison table claims "Fastest" for jax-js and "Fastest" for onnxruntime-web, which is a self-assessment in the project's own README, not a number you can plan capacity against. Treat the benchmark suite at jax-js.com/bench as the place to check, and run it on your own hardware.

## jax-js compared with TensorFlow.js and onnxruntime-web

The three runtimes in the README's comparison table differ in what they assume about your model. TensorFlow.js uses a TensorFlow-like API and is marked with a 2024 latest release in the table, which is the project's own reading of that ecosystem rather than a statement about its maintenance. It supports gradients and WebGL, and the table marks WebGPU as a preview. It does not support `jvp()`, `jit()`, `vmap()`, Jacobian or Hessian, sorting, Float64 or Float16. So if your work is training a small network in the browser, TensorFlow.js covers it; if your work is differentiating through a function or vectorizing a batched computation, jax-js gives you the transformation primitives instead. onnxruntime-web takes the opposite approach: it executes static ONNX graphs, supports no gradients at all, and the table marks its API style as "Static ONNX graphs" with no eager array API. Its advantage is model coverage and a mature ONNX execution path, at the cost of a 24 MB Wasm payload alongside the 90 KB gzip JavaScript. jax-js sits between them: eager like TensorFlow.js, transformation-oriented like JAX, with partial ONNX import. Notably, jax-js is the only one of the three the table marks as supporting general einsum fully, sorting, Float64, mixed devices and reading safetensors.

## Maintenance, licence and what an upgrade costs

The repository is not archived. The last push was on 2026-09-10, and the most recent release listed is jax/v0.1.24 on 2026-08-31, preceded by 0.1.23 on 2026-08-17 and 0.1.22 on 2026-08-13. That is a steady release cadence on the 0.1.x line. The version number is the main upgrade consideration: at 0.1.x there is no stability promise in the version itself, and the compatibility table in FEATURES.md is the document that tells you which operators exist on which backend, so an upgrade should be checked against it rather than against the changelog alone. The repository is a pnpm workspace, with `packages/`, `src/`, `test/`, `bench/` and `website/` at the top level, and the package's own scripts include `pnpm build`, `pnpm test` and `pnpm check`, so building from source is a normal pnpm workflow if you need to patch something. The licence is MIT, declared both in the repository LICENSE file and in the `license` field of package.json. MIT is permissive: it allows commercial and closed-source use and requires preserving the copyright notice. This is not legal advice; if you are redistributing the package inside a product, have your own counsel confirm the notice requirements.

## Conclusion

Adopt jax-js if your numerical code must run client-side in a browser and you want JAX semantics: the package is MIT licensed, published as @jax-js/jax, and the README lists WebGPU support as absent on Node.js and Firefox for Android. Do not adopt it if you need BFloat16, packed Uint8, or a server-side GPU path, since the compatibility table marks those unsupported or partial. Before committing, verify your target browser against the platforms table and check the operator you depend on in FEATURES.md, because the README describes ONNX model support as partial.

## FAQ

### How do I install jax-js?

Install it with npm i @jax-js/jax, or import it directly from esm.sh in a module script tag if you do not want a bundler. The README describes the CDN form as the easiest way to start on a blank HTML page.

### Can jax-js run on Node.js or only in the browser?

It runs on Node.js, but only with the Wasm CPU backend. The platforms table marks WebGPU and WebGL as absent for Node.js, and WebGPU on Deno as asynchronous only.

### Does jax-js support autodiff and JIT like JAX?

Yes. The feature table lists gradients, Jacobian and Hessian, jvp() forward differentiation, jit() kernel fusion, vmap() auto-vectorization and graph capture as supported.

### What licence does jax-js use?

MIT, declared in the repository LICENSE file and in the license field of package.json. That permits commercial and closed-source use provided the copyright notice is preserved.

## Sources

- [ekzhang/jax-js on GitHub](https://github.com/ekzhang/jax-js)
- [License: MIT](https://github.com/ekzhang/jax-js/blob/main/LICENSE)
- [Project website](https://jax-js.com)
- [README](https://github.com/ekzhang/jax-js/blob/main/README.md)
- [Releases](https://github.com/ekzhang/jax-js/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/ekzhang-jax-js