基本介紹
發展歷程,主要功能,性能特點,運行環境,
發展歷程
JAX本身的作用域很窄，專注於高效的數組操作程式轉變圍繞JAX構建的是不斷發展的機器學習生態系統， 數值計算工具。
主要功能
JAX是一個面向數組的數值計算庫，具有自動微分和JIT編譯功能，可實現高性能機器學習研究。
性能特點
差異化：基於一致性的最佳化是ML的基礎。JAX原生地支持任意數值函式的正向和反向模式自動微分，通過函式轉換，如格拉德、hessian、jacfwd和jacrev。
向量化：在機器學習研究中，經常將單個函式套用於大量數據，例如計算一批數據的損失或評估每個示例的梯度以進行差異化私有學習。JAX通過vmap轉換提供自動矢量化，簡化了這種形式的編程。例如，研究人員在實現新算法時不需要推理。JAX還通過相關的pmap轉換支持大規模數據並行，優雅地分發對於單個加速器的記憶體來說太大的數據。
JIT編譯：XLA用於在GPU和Cloud TPU加速器上實時（JIT）編譯和執行JAX程式。JIT編譯與JAX的NumPy一致性API一起，使之前沒有高性能計算經驗的研究人員能夠輕鬆擴展到一個或多個加速器。
運行環境
Linux, x86_64 | Linux, aarch64 | Mac, x86_64 | Mac, aarch64 | Windows, x86_64 | Windows WSL2, x86_64 | |
|---|---|---|---|---|---|---|
CPU | yes | yes | yes | yes | yes | yes |
NVIDIA GPU | yes | yes | no | n/a | no | experimental |
Google Cloud TPU | yes | n/a | n/a | n/a | n/a | n/a |
AMD GPU | experimental | no | experimental | n/a | no | no |
Apple GPU | n/a | no | n/a | experimental | n/a | n/a |
Intel GPU | experimental | n/a | n/a | n/a | no | no |
