How the jax.jit() JIT compiler works in jax-js - Eric

For this, I need to make a deep learning compiler from scratch, and I want to keep it lightweight (e.g., JAX uses XLA as its compiler, which is 200 KLoC — too much bundle size for the web!). This is a note about the trickiest fundamental problem I’ve run into, and how I’m going about solving it. JAX is a great library. It takes the numerical computing properties of NumPy, shoves in GPU + Autograd, then packages it all up in a convenient API. By writing JAX in pure JS, using web APIs, we solve two problems: How to do numerical compute in the browser? Like taking the mean of some numbers, or applying an image filter. Lots of applications, (statistics, data science, classical ML, CV, etc.), but right now it’s pretty hard to do well. How do you run GPU compute in the browser? There are technologies like WebGPU if you want to write your own shaders, which is great if you’re making a video game. But this is tricky if you just want to do something simple. After all, a lot more people use PyTo

saved by

• Sudarsh K