You should check out https://jax-js.com/. It has parity with quite a large surface area of Jax, and compiles natively in the browser to wasm and webGPU, entirely written in JS.