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.

Oh, this is neat! I knew posting something that claimed a library didn't exist would inevitably surface one :)