jax (pypi): Differentiate, compile, and transform Numpy code.
Keywords: jax