JAX utilise XLA pour compiler et exécuter vos programmes NumPy sur des GPU et des TPU. La compilation s’effectue sous le capot par défaut, les appels de bibliothèque étant compilés et exécutés juste à temps. Mais JAX vous permet également de compiler juste à temps vos propres fonctions Python dans des noyaux optimisés pour XLA à l’aide d’une API à fonction unique, jit. La compilation et la différenciation automatique peuvent être composées arbitrairement, afin que vous puissiez exprimer des algorithmes sophistiqués et obtenir des performances maximales sans quitter Python. Vous pouvez même programmer plusieurs GPU ou cœurs TPU à la fois à l’aide de pmap et différencier l’ensemble.