The JIT + tensor model ends up being a really nice way to express light transport all in simple python, so I reimplemented some new research papers from SIGGRAPH like REstir PG and SZ and it just works. instead of complicated cpp its just a 200 LOC of python.
https://blog.evjang.com/2019/11/jaxpt.html is a better demo of how to render the Cornell Box on a TPU using differentiable path tracing.
Whoa.. the cursor rule I didn't know I needed!