JAX: High performance array computing
Contents
JAX: High performance array computing#
High performance array computing
JAX is a Python library for accelerator-oriented array computation and program transformation, designed for high-performance numerical computing and large-scale machine learning.
JAX provides a familiar NumPy-style API for ease of adoption by researchers and engineers.
JAX includes composable function transformations for compilation, batching, automatic differentiation, and parallelization.
The same code executes on multiple backends, including CPU, GPU, & TPU
If you’re looking to use JAX to train neural networks, check out theJAX AIStack!
Ecosystem#
JAX itself is narrowly-scoped and focuses on efficient array operations & programtransformations. Built around JAX is an evolving ecosystem of machine learning andnumerical computing tools; the following is just a small sample of what is out there:
Optimizers & solvers
Data loading
Probabilistic modeling
Many more JAX-based libraries have been developed; the community-runAwesome JAX pagemaintains an up-to-date list.