@jax-js/jax 0.0.3 → 0.0.5

package/README.md

@@ -2,8 +2,9 @@
 
 [Website](https://www.ekzhang.com/jax-js/) | [API Reference](https://www.ekzhang.com/jax-js/docs/)
 
-This 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.
+**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.
 
 ```bash
 npm i @jax-js/jax
@@ -12,6 +13,10 @@ 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
 
 You can use `jax-js` as an array API, just like NumPy.
@@ -24,7 +29,7 @@ const x = np.array([1, 2, 3]);
 const y = x.mul(4); // [4, 8, 12]
 ```
 
-It also lets you take derivatives like in JAX.
+It also lets you take derivatives with `grad` like in JAX (as well as `vmap`, `jit`).
 
 ```js
 import { grad, numpy as np } from "@jax-js/jax";
@@ -37,42 +42,107 @@ const xnorm = norm(x.ref); // 1^2 + 2^2 + 3^2 = 14
 const xgrad = grad(norm)(x); // [2, 4, 6]
 ```
 
-The default backend runs on CPU, but on [supported browsers](https://caniuse.com/webgpu),
-you can switch to GPU for maximum performance.
+The default backend runs on CPU, but on [supported browsers](https://caniuse.com/webgpu) including
+Chrome and iOS Safari, you can switch to GPU for better performance.
 
 ```js
-import { numpy as np, setDevice } from "@jax-js/jax";
+import { defaultDevice, init, numpy as np } from "@jax-js/jax";
+
+// Initialize the GPU backend.
+await init("webgpu");
 
 // Change the default backend to GPU.
-setDevice("webgpu");
+defaultDevice("webgpu");
 
 const x = np.ones([4096, 4096]);
 const y = np.dot(x.ref, x); // JIT-compiled into a matrix multiplication kernel
 ```
 
+Most common JAX APIs are supported. See the [compatibility table](./FEATURES.md) for a full
+breakdown of what features are available.
+
+### Web usage (CDN)
+
+If you want to use `jax-js` 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.
+
+```html
+<script type="module">
+  import { numpy as np } from "https://esm.sh/@jax-js/jax";
+</script>
+```
+
+### Performance
+
+We haven't spent a ton of time optimizing yet, but performance is generally pretty good. `jit` is
+very helpful for fusing operations together, and it's a feature only available on the web in jax-js.
+The default kernel-tuning heuristics get about 3000 GFLOP/s for matrix multiplication on an M4 Pro
+chip ([try it](https://www.ekzhang.com/jax-js/bench/matmul)).
+
+For that example, it's around the same GFLOP/s as
+[TensorFlow.js](https://github.com/tensorflow/tfjs) and
+[ONNX Runtime Web](https://www.npmjs.com/package/onnxruntime-web), which both use handwritten
+libraries of custom kernels (versus jax-js, which generates kernels with an ML compiler).
+
+## Examples
+
+If you make something cool with jax-js, don't be a stranger! We can feature it here.
+
+- [In-browser REPL](https://www.ekzhang.com/jax-js/repl)
+- [Interactive MNIST training](https://www.ekzhang.com/jax-js/mnist)
+- [Matmul benchmark](https://www.ekzhang.com/jax-js/bench/matmul)
+- [Conv2d benchmark](https://www.ekzhang.com/jax-js/bench/conv2d)
+- [Mandelbrot set](https://www.ekzhang.com/jax-js/mandelbrot)
+
 ## Development
 
-Under construction.
+_The following technical details are for contributing to jax-js and modifying its internals._
+
+This repository is managed by [`pnpm`](https://pnpm.io/). You can compile and build all packages in
+watch mode with:
 
 ```bash
 pnpm install
 pnpm run build:watch
+```
+
+Then you can run tests in a headless browser using [Vitest](https://vitest.dev/).
 
-# Run tests
+```bash
 pnpm exec playwright install
 pnpm test
 ```
 
+We are currently on an older version of Playwright that supports using WebGPU in headless mode;
+newer versions skip the WebGPU tests.
+
+To start a Vite dev server running the website, demos and REPL:
+
+```bash
+pnpm -C website dev
+```
+
+## Future work / help wanted
+
+Contributions are welcomed in the following areas:
+
+- Adding support for more JAX functions and operations, see [compatibility table](./FEATURES.md).
+- Improving performance of the WebGPU and Wasm runtimes, generating better kernels, and using SIMD
+  and multithreading.
+- Helping the JIT compiler to fuse operations in more cases.
+- Adding WebGL runtime for older browsers that don't support WebGPU.
+- Making a fast transformer inference engine, comparing against onnxruntime-web.
+- Ergonomics and API improvements.
+
 ## Next on Eric's mind
 
 - Finish CLIP inference demo and associated features (depthwise convolution, vmap of gather, etc.)
-- Fix jit-of-grad returning very incorrect result
-- Improve perf of MNIST neural network
-  - Optimize conv2d further (maybe blocks -> local dims?)
-  - Add fused epilogue to JIT
-  - Reduce kernel overhead of constants / inline expressions
-- Investigate why jax-js Matmul is 2x slower on Safari TP than unroll kernel
-- How many threads to create per workgroup, depends on hardware
+- Performance
+  - Improve perf of MobileCLIP neural network
+    - Add fused epilogue to JIT
+    - Fix fusion of activation functions with branches like tanh
+    - Reduce kernel overhead of constants / inline expressions
+  - How many threads to create per workgroup, depends on hardware
 
 ## Milestones
 
@@ -91,9 +161,9 @@ pnpm test
 - [x] Other dtypes like int32 and bool
 - [x] `jit()` support via Jaxprs and kernel fusion
 - [x] We figure out the `dispose()` / refcount / linear types stuff
-  - [ ] `dispose()` for saved "const" tracers in Jaxprs
-  - [ ] Garbage collection for JIT programs
-  - [ ] Memory scheduling, buffer allocation (can be tricky)
+  - [x] `dispose()` for saved "const" tracers in Jaxprs
+  - [x] Garbage collection for JIT programs
+  - [x] Debug grad-grad-jit test producing a UseAfterFreeError
 - [ ] Demos: Navier-Stokes, neural networks, statistics
 - [x] Features for neural networks
   - [x] Convolution
@@ -103,6 +173,10 @@ pnpm test
   - [x] Better memory allocation that frees buffers
   - [ ] SIMD support for Wasm backend
   - [ ] Async / multithreading Wasm support
-- [ ] Device switching with `.to()` between webgpu/cpu/wasm
-- [ ] numpy/jax API compatibility table
-- [ ] Import tfjs models
+- [ ] Full support of weak types and committed devices
+  - [x] High-level ops have automatic type promotion
+  - [x] Weak types - [ref](https://docs.jax.dev/en/latest/type_promotion.html#weak-types)
+  - [ ] Committed devices -
+        [ref](https://docs.jax.dev/en/latest/sharded-computation.html#sharded-data-placement)
+  - [ ] Device switching with `device_put()` between webgpu/cpu/wasm
+- [x] numpy/jax API compatibility table