The author's workload is sort of different than the usual ML workload since the author's expression tree is large (10k nodes), while a modern neural net has a r

The author's workload is sort of different than the usual ML workload since the author's expression tree is large (10k nodes), while a modern neural net has a relatively smaller expression tree, maybe fewer than 100 for the larger neural nets? Another commenter mentioned Dr.Jit which seems to be designed for this use case. This is a quote from their project page. > Why did we create Dr.Jit, when dynamic derivative compilation is already possible using Python-based ML frameworks like JAX, Tensorflow, and PyTorch along with backends like XLA and TorchScript? > The reason is related to the typical workloads: machine learning involves small-ish computation graphs that are, however, made of arithmetically intense operations like convolutions, matrix multiplications, etc. The application motivating Dr.Jit (differentiable rendering) creates giant and messy computation graphs consisting of 100K to millions of "trivial" nodes (elementary arithmetic operations). In our experience, ML compilation backends use internal representations and optimization passes that are too rich for this type of input, causing them to crash or time out during compilation. If you have encountered such issues, you may find Dr.Jit useful.