Speeding up my Ray Tracer using JAX (kayleegeorge.github.io)

The author rewrote the "Ray Tracing in One Weekend" tutorial from Rust into JAX to turn a painfully slow CPU implementation into a fast, GPU-friendly renderer with minimal extra engineering. Ray tracing maps well to JAX: pixel computations are embarrassingly parallel, heavy on vector math, and benefit from differentiation (useful for inverse rendering and neural radiance fields). By expressing the tracer as pure functions and arrays rather than objects, the author leverages JAX’s jit and vmap to compile and run thousands of rays in parallel and to reuse a single compiled kernel per image resolution. Key technical moves that enabled the speedups: represent rays and scene data as jnp arrays (no classes/side-effects), avoid runtime control-flow early returns by computing quantities for all branches and using jnp.where masking so JAX can trace and differentiate, and mark invariants (image width/height) as static_argnums so JIT doesn’t recompile per call. Vectorize over pixels with vmap and over spheres simultaneously to test all intersections in parallel, then pick the nearest hit via argmin on masked t-values. The guide also covers camera construction recomputed per call (JIT eliminates overhead), handling near/far roots for refraction, and the need to use JAX’s stateless PRNG model for scattering. The result is a compact, differentiable, and highly parallel ray tracer (runnable in the author’s Colab) that scales from single spheres to complex scenes with substantial speed gains.