Write NumPy code in the familiar way, and then run it on GPUs and TPUs. Automatic differentiation and data batching will be handled automatically, without requiring any extra effort from you.

- Automatically compute derivatives of standard Python and NumPy functions.
- Compile pure functions using jit to run them on hardware accelerators.
- Combine grad, jit, and vmap in any order and scale your computations across multiple devices.

https://github.com/jax-ml/jax

https://t.iss.one/MachineLearning9