Google Jax - ไลบรารีแมชชีนเลิร์นนิงประสิทธิภาพสูง

"สร้างสิ่งที่ใช้งานง่ายให้ทำงานได้รวดเร็วแล้วนำไปใช้กับแมชชีนเลิร์นนิง"

• ผสานรวมเพียง Python และ Numpy

→ ใช้ XLA เพื่อคอมไพล์และรัน Numpy บน GPU/TPU

→ สามารถทำ JIT compile ฟังก์ชัน Python ด้วย API เดียว และใส่ลงในเคอร์เนลที่ XLA ปรับแต่งไว้ได้อย่างง่ายดาย

→ รันบน GPU/TPU หลายตัวได้ง่ายเช่นกัน (vmap, pmap)

• ประสิทธิภาพเหนือกว่า Python+Numpy แบบเดิมอย่างมาก

[xguru]

DeepMind ได้รีแฟกเตอร์ทั้งหมดโดยอิงกับ Jax

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