How DeepMind Is Using JAX To Accelerate AI Research

JAX is a Python library developed by Google researchers for high-performance numerical computing. Its API is based on NumPy. NumPy is a collection of functions applied in scientific computing. Developers extensively adopt Python and NumPy, making JAX simple, flexible, and easy to use. JAX and its developing ecosystem of open source libraries have assisted and accelerated numerous machine learning projects.