JAX

NumPy-like array library with autograd and XLA compilation.

JAX from Google Research extends NumPy with automatic differentiation and just-in-time compilation to CPU, GPU, and TPU via XLA. It is widely used for high-performance machine-learning research and scientific computing.

License: Apache-2.0

Category: AI / ML

Website: https://github.com/google/jax

Install

pip install jax

PreviousJaeger NextJekyll