@jax-js/jax 0.0.1 → 0.0.2

package/README.md

@@ -55,18 +55,23 @@ 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
+- 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
-  - 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
 
@@ -88,12 +93,12 @@ 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
+- [x] Features for neural networks
   - [ ] Convolution
-  - [ ] Random and initializers
+  - [x] Random and initializers
   - [ ] Optimizers (optax package?)
 - [ ] Wasm backend (needs malloc)
   - [ ] SIMD support for Wasm backend