"Membuat hal yang mudah digunakan menjadi cepat lalu menerapkannya pada machine learning"
- Hanya menggabungkan Python dan Numpy
→ Mengompilasi dan menjalankan Numpy di GPU/TPU menggunakan XLA
→ Fungsi Python dapat di-JIT compile hanya dengan satu API sehingga mudah dimasukkan ke kernel yang dioptimalkan XLA
→ Eksekusi di banyak GPU/TPU juga mudah (vmap, pmap)
- Jauh melampaui performa Python+Numpy yang ada
1 komentar
DeepMind telah merefaktorisasi seluruh sistem berbasis Jax
https://deepmind.com/blog/article/using-jax-to-accelerate-our-research