Google Jax - thư viện machine learning hiệu năng cao

"Tạo ra những gì dễ dùng một cách nhanh chóng và áp dụng vào machine learning"

• Chỉ kết hợp Python và Numpy

→ Dùng XLA để biên dịch và chạy Numpy trên GPU/TPU

→ Có thể JIT compile hàm Python chỉ với một API và dễ dàng đưa vào các kernel được XLA tối ưu hóa

→ Cũng dễ dàng chạy trên nhiều GPU/TPU (vmap, pmap)

• Hiệu năng vượt xa Python+Numpy truyền thống

[xguru]

DeepMind đã refactor toàn bộ dựa trên Jax

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