GoKawiilJAX’s true calling: Ray-Marching renderers in Python on WebGL
Demo
(move your mouse/thumb across the image)
Why, though?
Well, I’ve been drooling over this tool the cool kids use, and wondering how I can join the gang. It’s called JAX.
It’s got GPU accelerated functions over n-dimensional arrays. And built-in compile-time differentiability of these!? Auto-vectorization?? And you just have to do like with numpy. What’s not to like? Go home APL! So I’ve been doing the obvious, the thing JAX was truly meant for: a graphics renderer.
Why, do you ask? Well, the animated image above is a 3-dimensional [512 pixels][512 pixels][3 colors] array for starters (or tensor, if you like). And we can define its content from the output of a function. Start from mouse position and time input, plug in some maths, hard-code a sphere and a cube in there, and voilà, pixels are painted!
And for our first trick here is the code, at about just 100 lines of Python. Yes, Python for browser code, because JAX can also be exported and run on the browser. On WebGL.
Below are some of the techniques used and where JAX shines:
Distances
... continue reading