Google Jax - biblioteca de machine learning de alto desempenho

"Tornar fácil de usar algo rápido e aplicá-lo ao machine learning"

• Combina apenas Python e Numpy

→ Compila e executa o Numpy em GPU/TPU usando XLA

→ É possível fazer JIT compile de funções Python com uma única API e inseri-las facilmente em kernels otimizados por XLA

→ Também facilita a execução em múltiplas GPU/TPU (vmap, pmap)

• Supera em muito o desempenho do Python+Numpy tradicional

[xguru]

A DeepMind refatorou tudo com base em Jax

https://deepmind.com/blog/article/using-jax-to-accelerate-our-research