Communication Dans Un Congrès Année : 2019

Compiling Machine Learning Programs via High-Level Tracing

Roy Frostig
  • Fonction : Auteur
  • PersonId : 1567515
Chris Leary
  • Fonction : Auteur
  • PersonId : 1567516

Résumé

We describe JAX, a domain-specific tracing JIT compiler for generating high-performance accelerator code from pure Python and Numpy machine learning programs. JAX uses the XLA compiler infrastructure to generate optimized code for the program subroutines that are most favorable for acceleration, and these optimized subroutines can be called and orchestrated by arbitrary Python. Because the system is fully compatible with Autograd, it allows forward-and reverse-mode automatic differentiation of Python functions to arbitrary order. Because JAX supports structured control flow, it can generate code for sophisticated machine learning algorithms while maintaining high performance. We show that by combining JAX with Autograd and Numpy we get an easily programmable and highly performant ML system that targets CPUs, GPUs, and TPUs, capable of scaling to multi-core Cloud TPUs.

Fichier principal
Vignette du fichier
jax-mlsys2018.pdf (618.89 Ko) Télécharger le fichier
Origine Fichiers produits par l'(les) auteur(s)
Licence

Dates et versions

hal-05188750 , version 1 (27-07-2025)

Licence

Identifiants

  • HAL Id : hal-05188750 , version 1

Citer

Roy Frostig, Matthew James Johnson, Chris Leary. Compiling Machine Learning Programs via High-Level Tracing. SysML Conference 2018, Mar 2019, Stanford, United States. ⟨hal-05188750⟩
289 Consultations
328 Téléchargements

Partager

  • More