@jax-js/jax 0.0.2 → 0.0.4

package/README.md

@@ -1,6 +1,6 @@
 # jax-js: JAX in pure JavaScript
 
-[Website](https://www.ekzhang.com/jax-js/)
+[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.
@@ -37,43 +37,58 @@ 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), you can
+switch to GPU for better performance.
 
 ```js
-import { numpy as np, setDevice } from "@jax-js/jax";
+import { defaultDevice, numpy as np } from "@jax-js/jax";
 
 // 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.
+
 ## Development
 
-Under construction.
+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 seem to skip the WebGPU tests._
+
+To start a Vite dev server running the website, demos and REPL:
+
+```bash
+pnpm -C website dev
+```
+
 ## Next on Eric's mind
 
-- Fix jit-of-grad returning very incorrect result
-- Probably add static_argnums to jit() so that clip and some nn functions have jit added
-- Improve perf of MNIST neural network
-  - Adding fused reductions 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
-- Think about two-stage `cumsum()`
-- Frontend transformations need to match backend type for pureArray() and zeros() calls
+- Finish CLIP inference demo and associated features (depthwise convolution, vmap of gather, etc.)
+- Performance
+  - 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
 
 ## Milestones
 
@@ -92,16 +107,33 @@ 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
-  - [ ] Convolution
+  - [x] Convolution
   - [x] Random and initializers
-  - [ ] Optimizers (optax package?)
-- [ ] Wasm backend (needs malloc)
+  - [x] Optimizers (optax package?)
+- [x] Wasm backend (needs malloc)
+  - [x] Better memory allocation that frees buffers
   - [ ] SIMD support for Wasm backend
-- [ ] Device switching with `.to()` between webgpu/cpu/wasm
-- [ ] numpy/jax API compatibility table
-- [ ] Import tfjs models
+  - [ ] Async / multithreading Wasm support
+- [ ] Full support of weak types and committed devices
+  - [ ] High-level ops have automatic type promotion
+  - [ ] 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
+
+## 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, using SIMD and
+  multithreading.
+- 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.