JAX is an open-source numerical computing and machine-learning library from Google.
AI researchers and engineers use it for automatic differentiation, accelerated array computation, model research, and high-performance machine learning.