@jax-js/jax 0.1.9 → 0.1.11
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/README.md +35 -19
- package/dist/{backend-BId79r5b.js → backend-DZvR7mZV.js} +831 -26
- package/dist/{backend-DpI0riom.cjs → backend-DlYlOYqN.cjs} +872 -25
- package/dist/index.cjs +364 -20
- package/dist/index.d.cts +175 -11
- package/dist/index.d.ts +175 -11
- package/dist/index.js +363 -21
- package/dist/{webgl-DnGrclTz.js → webgl-D8-14NzA.js} +7 -1
- package/dist/{webgl-C5NjXc1p.cjs → webgl-Ovaaa-Qx.cjs} +7 -1
- package/dist/{webgpu-AN0cG_nB.js → webgpu-Dg8FpYrH.js} +141 -6
- package/dist/{webgpu-CdjiJSa7.cjs → webgpu-uU9nnttc.cjs} +141 -6
- package/package.json +5 -16
package/README.md
CHANGED
|
@@ -43,18 +43,37 @@ way to get started on a blank HTML page.
|
|
|
43
43
|
</script>
|
|
44
44
|
```
|
|
45
45
|
|
|
46
|
+
### Platforms
|
|
47
|
+
|
|
48
|
+
This table refers to latest versions of each browser. WebGPU has gained wide support in browsers as
|
|
49
|
+
of late 2025.
|
|
50
|
+
|
|
51
|
+
| Platform | CPU (Wasm) | GPU (WebGPU) | GPU (WebGL) |
|
|
52
|
+
| ------------------- | ---------- | -------------- | ----------- |
|
|
53
|
+
| Chrome / Edge | ✅ | ✅ | ✅ |
|
|
54
|
+
| Firefox | ✅ | ✅ - macOS 26+ | ✅ |
|
|
55
|
+
| Safari | ✅ | ✅ - macOS 26+ | ✅ |
|
|
56
|
+
| iOS | ✅ | ✅ - iOS 26+ | ✅ |
|
|
57
|
+
| Chrome for Android | ✅ | ✅ | ✅ |
|
|
58
|
+
| Firefox for Android | ✅ | ❌ | ✅ |
|
|
59
|
+
| Node.js | ✅ | ❌ | ❌ |
|
|
60
|
+
| Deno | ✅ | ✅ - async | ❌ |
|
|
61
|
+
|
|
46
62
|
## Examples
|
|
47
63
|
|
|
48
|
-
|
|
64
|
+
Community usage:
|
|
49
65
|
|
|
66
|
+
- [**autoresearch-webgpu**: autoresesarch, in the browser](https://autoresearch.lucasgelfond.online/)
|
|
50
67
|
- [**tanh.xyz**: Interactive ML visualizations](https://tanh.xyz/)
|
|
68
|
+
- [**jax-js-bayes**: Declarative Bayesian modeling library](https://github.com/StefanSko/jax-js-bayes)
|
|
51
69
|
|
|
52
|
-
|
|
70
|
+
Demos on the jax-js website:
|
|
53
71
|
|
|
54
72
|
- [Training neural networks on MNIST](https://jax-js.com/mnist)
|
|
55
73
|
- [Voice cloning: Kyutai Pocket TTS](https://jax-js.com/tts)
|
|
56
74
|
- [CLIP embeddings for books in-browser](https://jax-js.com/mobileclip)
|
|
57
75
|
- [Object detection: DETR ResNet-50 (ONNX)](https://jax-js.com/detr-resnet-50)
|
|
76
|
+
- [Fluid simulation (Navier-Stokes)](https://jax-js.com/fluid-sim)
|
|
58
77
|
- [In-browser REPL](https://jax-js.com/repl)
|
|
59
78
|
- [Matmul benchmark](https://jax-js.com/bench/matmul)
|
|
60
79
|
- [Conv2d benchmark](https://jax-js.com/bench/conv2d)
|
|
@@ -292,7 +311,9 @@ and determines how to execute compiled operations on them.
|
|
|
292
311
|
There are currently 4 devices in jax-js:
|
|
293
312
|
|
|
294
313
|
- `cpu`: Slow, interpreted JS, only meant for debugging.
|
|
295
|
-
- `wasm`: [WebAssembly](https://webassembly.org/),
|
|
314
|
+
- `wasm`: [WebAssembly](https://webassembly.org/), multi-threaded when
|
|
315
|
+
[`SharedArrayBuffer`](https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/SharedArrayBuffer)
|
|
316
|
+
is available.
|
|
296
317
|
- `webgpu`: [WebGPU](https://developer.mozilla.org/en-US/docs/Web/API/WebGPU_API), available on
|
|
297
318
|
[supported browsers](https://caniuse.com/webgpu) (Chrome, Firefox, Safari, iOS).
|
|
298
319
|
- `webgl`: [WebGL2](https://developer.mozilla.org/en-US/docs/Web/API/WebGL2RenderingContext), via
|
|
@@ -337,6 +358,8 @@ self-contained way in other projects.
|
|
|
337
358
|
|
|
338
359
|
### Performance
|
|
339
360
|
|
|
361
|
+
To see per-kernel traces in browser development tools, call `jax.profiler.startTrace()`.
|
|
362
|
+
|
|
340
363
|
The WebGPU runtime includes an ML compiler with tile-aware optimizations, tuned for indiidual
|
|
341
364
|
browsers. Also, this library uniquely has the `jit()` feature that fuses operations together and
|
|
342
365
|
records an execution graph. jax-js achieves **over 7000 GFLOP/s** for matrix multiplication on an
|
|
@@ -367,19 +390,6 @@ pnpm install
|
|
|
367
390
|
pnpm run build:watch
|
|
368
391
|
```
|
|
369
392
|
|
|
370
|
-
The `pnpm install` command automatically sets up Git hooks via
|
|
371
|
-
[Husky](https://typicode.github.io/husky/). Pre-commit hooks will run ESLint and Prettier on staged
|
|
372
|
-
files to ensure code quality.
|
|
373
|
-
|
|
374
|
-
You can also run linting and formatting manually:
|
|
375
|
-
|
|
376
|
-
```bash
|
|
377
|
-
pnpm lint # Run ESLint
|
|
378
|
-
pnpm format # Format all files with Prettier
|
|
379
|
-
pnpm format:check # Check formatting without writing
|
|
380
|
-
pnpm check # Run TypeScript type checking
|
|
381
|
-
```
|
|
382
|
-
|
|
383
393
|
Then you can run tests in a headless browser using [Vitest](https://vitest.dev/).
|
|
384
394
|
|
|
385
395
|
```bash
|
|
@@ -396,6 +406,15 @@ To start a Vite dev server running the website, demos and REPL:
|
|
|
396
406
|
pnpm -C website dev
|
|
397
407
|
```
|
|
398
408
|
|
|
409
|
+
You can run the linter, code formatter, and type checker with:
|
|
410
|
+
|
|
411
|
+
```bash
|
|
412
|
+
pnpm lint # Run ESLint
|
|
413
|
+
pnpm format # Format all files with Prettier
|
|
414
|
+
pnpm format:check # Check formatting without writing
|
|
415
|
+
pnpm check # Run TypeScript type checking
|
|
416
|
+
```
|
|
417
|
+
|
|
399
418
|
## Future work / help wanted
|
|
400
419
|
|
|
401
420
|
Contributions are welcomed! Some fruitful areas to look into:
|
|
@@ -403,9 +422,6 @@ Contributions are welcomed! Some fruitful areas to look into:
|
|
|
403
422
|
- Adding support for more JAX functions and operations, see [compatibility table](./FEATURES.md).
|
|
404
423
|
- Improving performance of the WebGPU and Wasm runtimes, generating better kernels, and using SIMD
|
|
405
424
|
and multithreading. (Even single-threaded Wasm could be ~20x faster.)
|
|
406
|
-
- Adding support for `jax.profiling`, in particular the start and end trace functions. We should be
|
|
407
|
-
able to generate `traceEvents` from backends (especially on GPU, with precise timestamp queries)
|
|
408
|
-
to help with model performance debugging.
|
|
409
425
|
- Helping the JIT compiler to fuse operations in more cases, like `tanh` branches.
|
|
410
426
|
- Making a fast transformer inference engine, comparing against onnxruntime-web.
|
|
411
427
|
|