@jax-js/jax 0.0.1

package/LICENSE

@@ -0,0 +1,20 @@
+MIT License
+
+Copyright (c) 2025 Eric Zhang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of
+this software and associated documentation files (the "Software"), to deal in
+the Software without restriction, including without limitation the rights to
+use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
+the Software, and to permit persons to whom the Software is furnished to do so,
+subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
+FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
+COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
+IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
+CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

package/README.md

@@ -0,0 +1,102 @@
+# jax-js: JAX in pure JavaScript
+
+[Website](https://www.ekzhang.com/jax-js/)
+
+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.
+
+```bash
+npm i @jax-js/jax
+```
+
+Under the hood, it translates array operations into a compiler representation, then synthesizes
+kernels in WebAssembly and WebGPU.
+
+## Quickstart
+
+You can use `jax-js` as an array API, just like NumPy.
+
+```js
+import { numpy as np } from "@jax-js/jax";
+
+// Array operations, compatible with NumPy.
+const x = np.array([1, 2, 3]);
+const y = x.mul(4); // [4, 8, 12]
+```
+
+It also lets you take derivatives like in JAX.
+
+```js
+import { grad, numpy as np } from "@jax-js/jax";
+
+// Calculate derivatives with reverse-mode AD.
+const norm = (a) => a.ref.mul(a).sum();
+
+const x = np.array([1, 2, 3]);
+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.
+
+```js
+import { numpy as np, setDevice } from "@jax-js/jax";
+
+// Change the default backend to GPU.
+setDevice("webgpu");
+
+const x = np.ones([4096, 4096]);
+const y = np.dot(x.ref, x); // JIT-compiled into a matrix multiplication kernel
+```
+
+## Development
+
+Under construction.
+
+```bash
+npm install
+npm run build:watch
+npm test
+```
+
+## Next on Eric's mind
+
+- softmax, log_softmax, logsumpexp, one_hot
+- Start working on first neural network
+- 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
+
+- [x] It works!
+- [x] Demos: Browser REPL / editor
+- [x] First custom kernel
+- [x] Custom WebGPU backend, removing tfjs dependency
+  - [x] Low-level operations
+  - [x] Create `class Array {}` wrappers
+  - [x] Reduction operations
+- [ ] Kernel tuning (see `tuner.ts`)
+  - [x] "Upcast" optimizations (compute a tile per thread, e.g., matmul)
+  - [x] "Unroll" optimizations (multiple loop iters per thread, e.g., matmul)
+  - [ ] "Group" optimizations (multiple threads per value, e.g., matvec)
+  - [ ] Blocks respect local dimensions
+- [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, maybe needs to be moved off-device
+  - [ ] 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)
+  - [ ] SIMD support for Wasm backend
+- [ ] Device switching with `.to()` between webgpu/cpu/wasm
+- [ ] numpy/jax API compatibility table
+- [ ] Import tfjs models