JAX is a sophisticated and well implemented high level tracer, combined with a backend compiler that makes use of XLA. High level tracing has some fundamental draw backs in key features we want (efficient scalar AD for example), so the approach being advocated here is doing AD as a compiler transform on the original source language. Mike has some details in an earlier paper: https://arxiv.org/abs/1810.07951. As a shameless plug, you can plug (heh) Julia into XLA as well, e.g. for targeting TPUs (https://arxiv.org/abs/1810.09868). In fact, Julia and JAX are probably better at that than TensorFlow, because the TF execution model is a bit of a mismatch for what XLA expects.