JAX

lastmod 2024-08-21created 2024-08-211 edits

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

Google JAX is a machine learning framework for transforming numerical functions. It is described as combining a modified version of autograd with TensorFlow's XLA.

See also

• XLA
• PyTorch

Favorite site

• Github - google/jax
• High-performance deep learning frameworks: JAX/Flax - Modulabs