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.
JAX at DeepMind
While supporting AI research, it is essential to ensure that the AI experiments are scalable to the real-world application. Advancing AI research needs balancing rapid prototyping and quick iteration. Researchers at DeepMind feature approaches to enable the core JAX libraries to continue with new research directions.