JAX

JAX is a library for array-oriented numerical computation with automatic differentiation and JIT to enable high performance machine learning research

https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html

Thinking in JAX §