Jax's true calling: Ray-Marching renderers on WebGL

JAX’s true calling: Ray-Marching renderers in Python on WebGL

Why, though?

JAX is widely known for GPU accelerated functions over n-dimensional arrays, built-in compile-time differentiability, and auto-vectorization. While typically used for Machine Learning, it is also an incredible tool for a graphics renderer.

An animated 3D image is essentially a 3-dimensional array (or tensor). By defining its content from the output of a function—starting from mouse position and time input, then applying math—pixels are painted. With JAX, this can be achieved in about 100 lines of Python code, exported to run on the browser via WebGL.

Distances and SDFs

Instead of drawing polygons, this approach uses Signed Distance Functions (SDF). These functions define an object by the distance to its surface (negative inside, positive outside).

• Composability: You can create unions of objects by taking the minimum of their SDFs, or intersections by taking the maximum.
• Ray-marching: If the closest point to an object is a length L away, you can move by length L in any direction without hitting it. This is the core of the sphere tracing algorithm.

Vectorization and Gradients

JAX shines in its ability to transform functions:

• vmap: Using vmap, we can transform functions assigned to single pixels into functions that compute all pixels of an image in parallel.
• Automatic Differentiation: The gradient of the SDF is the normal to the surface. With JAX's jax.grad, light diffusion and reflection can be calculated at compile-time with mathematical purity, avoiding run-time epsilon tricks.

Mechanical Sympathy for Maths

JAX allows the code to be 70% math, acting as a high-level DSL that stays very close to the metal. It bridges the gap between complex mathematical expressions and high-performance GPU execution.