Exploring JAX for Graphics Rendering with Signed Distance Functions

I've been captivated by JAX, a tool that offers GPU-accelerated functions over n-dimensional arrays and compile-time differentiability. It allows for auto-vectorization, making it a powerful tool for graphics rendering. By using JAX, I can create a 3D animated image as a [512 pixels][512 pixels][3 colors] array, defining its content through a function. Starting with inputs like mouse position and time, and incorporating mathematical functions, I can render images with just about 100 lines of Python code that run on the browser via WebGL.

## Distances and Signed Distance Functions

Instead of drawing polygons, I use Signed Distance Functions (SDFs), which are elegant and versatile. A sphere, for instance, is simply a distance to a point, and a cylinder is a distance to a line. SDFs are composable, allowing for operations like unions and intersections of objects using functions like min() and max(). For smooth unions, we preserve differentiability, and for intersections, a softmax function can be used.

SDFs also facilitate movement in space without collision, thanks to the raymarching or sphere tracing algorithm. With JAX's vmap, I can vectorize these functions, transforming single-pixel operations into parallel computations across all pixels.

## Differentiability and Mathematical Elegance

The gradient of an SDF provides the normal to a surface, enabling light diffusion or reflection at compile-time with just one line of code. This eliminates the need for run-time adjustments, achieving mathematical purity and efficiency. Although JAX wasn't created solely for my convenience, its primitives allow for code that is predominantly mathematical yet close to the hardware.

## Future Explorations

I am eager to explore further possibilities with JAX, such as using jax-js for WebGPU support, experimenting with JAX Geometric Algebra for expressing SDFs in Projective Geometry, and visualizing the relationship between Neural Networks and 3D/SDFs. Additionally, jax.experimental.jet could potentially enhance the ray-marching process with Taylor expansions of functions.