@jax-js/jax 0.0.1 → 0.0.3

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.
@@ -55,20 +55,24 @@ const y = np.dot(x.ref, x); // JIT-compiled into a matrix multiplication kernel
 Under construction.
 
 ```bash
-npm install
-npm run build:watch
-npm test
+pnpm install
+pnpm run build:watch
+
+# Run tests
+pnpm exec playwright install
+pnpm test
 ```
 
 ## Next on Eric's mind
 
-- softmax, log_softmax, logsumpexp, one_hot
-- Start working on first neural network
+- 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
-  - Need to break up kernel dispatches if workgroup count exceeds 65536
-- Think about two-stage `cumsum()`
-- Frontend transformations need to match backend type for pureArray() and zeros() calls
 
 ## Milestones
 
@@ -88,15 +92,17 @@ npm test
 - [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, maybe needs to be moved off-device
+  - [ ] Garbage collection for JIT programs
   - [ ] Memory scheduling, buffer allocation (can be tricky)
 - [ ] Demos: Navier-Stokes, neural networks, statistics
-- [ ] Features for neural networks
-  - [ ] Convolution
-  - [ ] Random and initializers
-  - [ ] Optimizers (optax package?)
-- [ ] Wasm backend (needs malloc)
+- [x] Features for neural networks
+  - [x] Convolution
+  - [x] Random and initializers
+  - [x] Optimizers (optax package?)
+- [x] Wasm backend (needs malloc)
+  - [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