Accelerating Machine Learning using JAX · Luma

JAX is a system for high-performance machine-learning research. It offers the familiarity of Python NumPy together with hardware acceleration. JAX enables the definition and composition of user-wielded function transformations useful for machine-learning programs. These transformations include automatic differentiation, automatic batching, end-to-end compilation (via XLA), parallelizing over multiple accelerators, and more. Composing these transformations is the key to JAX's power and simplicity.