JAX allows to write auto-differentiable functions. It provides a NumPy and native Python compatible interface built on composable function transformations. Further optimization happens by automatic vectorization and running code on GPUs/TPUs.
Documentation: https://jax.readthedocs.io
Project repo: https://github.com/google/jax
