StochasticAD.jl

structs that void lots of optimizations which have to then be worked around, which is why Python JITs choose subsets of the language to avoid specific parts that are not easy to optimize or not possible to optimize. This is why each take a domain-specific subset, a different subset of the language for numba vs jax vs etc., to choose something that is nice for ML vs for more generic codes. With all of that, it's perfectly reasonable to point out that there have been languages which have been designed to not have the compilation difficulties, which have resulted having a single (JIT) compiler for the language. And by extension, it has made building machine learning and autodiff libraries not something that's a Google or Meta scale project (for example, PyTorch involves building GPU code bindings and a specialized JIT, not something very accessible). Julia is a language to point to here, but I think well-designed static languages like Rust also deserve a mention. How much further would we have gone if every new ML project didn't build a new compiler and a new automatic differentiation engine? What if the development was more modular and people could easy just work on the one thing they cared about? As a nice example, for last NeurIPS we put out a paper on automatic differentiation of discrete stochastic models, i.e. extending AD to automatically handle cases like agent-based models. The code is open source (https://github.com/gaurav-arya/StochasticAD.jl), and you can see it's almost all written by a (talented) undergraduate over a span of about 6 months. It requires the JIT compilation because it works on a lot of things that are not solely in big matrix multiplication GPU kernels, but Julia provides that. And multiple dispatch gives GPU support. Done. The closest thing in PyTorch, storchastic, gets exponential scaling instead of StochasticAD's linear, and isn't quite compatible with a lot of what's required for ML, so it benchmarks as thousands of times slower than the simple Julia code. Of course, when Meta needs it they can and will put the minds of 5-10 top PhDs on it to build it out into a feature of PyTorch over 2 years and have a nice release. But at the end of the day we really need to ask, is that how it should be?