JAX 結(jié)合了 Autograd 和 XLA，專門用于高性能機(jī)器學(xué)習(xí)研究。

　　憑借 Autograd，JAX 可以求導(dǎo)循環(huán)、分支、遞歸和閉包函數(shù)，并且它可以進(jìn)行三階求導(dǎo)。通過 grad，它支持自動(dòng)模式反向求導(dǎo)(反向傳播)和正向求導(dǎo)，且二者可以任何順序任意組合。

　　得力于 XLA，可以在 GPU 和 TPU 上編譯和運(yùn)行 NumPy 程序。默認(rèn)情況下，編譯發(fā)生在底層，庫調(diào)用實(shí)時(shí)編譯和執(zhí)行。但是 JAX 還允許使用單一函數(shù) API jit 將 Python 函數(shù)及時(shí)編譯為 XLA 優(yōu)化的內(nèi)核。編譯和自動(dòng)求導(dǎo)可以任意組合，因此可以在 Python 環(huán)境下實(shí)現(xiàn)復(fù)雜的算法并獲得最大的性能。

　　更深入地看，JAX 實(shí)際上是一個(gè)可擴(kuò)展的可組合函數(shù)轉(zhuǎn)換系統(tǒng)，grad 和 jit 都是這種轉(zhuǎn)換的實(shí)例。
　　（邯鄲網(wǎng)站建設(shè)）