Games and simulations are typically stateful, I'm not sure the functional purity of JAX is a good fit.

Also, what's the story for JAX + WebGL when it comes to targeting hardware-accelerated ray tracing?

Maybe you mean mutation? State is inherently a part of functional purity. It's just handled explicitly instead of implicitly.