I want to know how Jax turns Python functions into Jaxprs when doing JIT and other operations. Some simple operations seem straight forward, but how does it compile more advanced NumPy functions like broadcast and reshape? It seems like I either get very high level tutorials or have to go into the source code. Is there any Post notation specification for how the compilation/translation is done from Python to Jaxpr? Any papers that explains this? Continue reading...
