In this paper, we present initial work to compile general Julia code to TPU using this interface. This approach is in contrast to the approach taken by TensorFlow [1]}, which does not compile Python code proper, but rather uses Python to build a computational graph, which is then compiled. It is aesthetically similar to JAX [2]}, which does aim to offload computations written in Python proper by tracing and offloading high-level array operations. Crucially, however, we do not rely on tracing, instead we leverage Julia's static analysis and compilation capabilities to compile the full program, including any control flow to the device. In particular, our approach allows users to take advantage of the full expressiveness of the Julia programming language in writing their models. This includes higher-level features such as multiple dispatch, higher order functions and existing libraries such as those for differential equation solvers [3]} and generic linear algebra routines. Since it operates on pure Julia code, it is also compatible with the Zygote.jl [4]} automatic differentiation tool, which performs automatic differentiation as a high-level compiler pass. Putting these together, we are able to compile full machine learning models written using the Flux machine learning framework, fusing the forward and backwards model passes as well as the training loop into a single executable that is offloaded to the TPU.
