Google Jax - biblioteca de aprendizaje automático de alto rendimiento

"Crear rápidamente cosas fáciles de usar y aplicarlas al aprendizaje automático"

• Combina solo Python y Numpy

→ Usa XLA para compilar y ejecutar Numpy en GPU/TPU

→ Las funciones de Python se pueden compilar con JIT con una sola API y llevar fácilmente a kernels optimizados con XLA

→ También facilita la ejecución en múltiples GPU/TPU (vmap, pmap)

• Supera ampliamente el rendimiento de la combinación tradicional de Python+Numpy

[xguru]

DeepMind refactorizó todo basándose en Jax.

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