„Einfach nutzbare Dinge schnell machen und auf Machine Learning anwenden“
- Nur die Kombination aus Python und Numpy
→ Mit XLA wird Numpy kompiliert und auf GPU/TPU ausgeführt
→ Python-Funktionen lassen sich mit einer einzigen API per JIT kompilieren und einfach in XLA-optimierte Kernel einbinden
→ Auch die Ausführung auf mehreren GPUs/TPUs ist einfach (vmap, pmap)
- Übertrifft die Performance von bestehendem Python+Numpy bei Weitem
1 Kommentare
DeepMind hat vollständig auf Jax-Basis refaktorisiert.
https://deepmind.com/blog/article/using-jax-to-accelerate-our-research