JAX

JAX，由Google开发、並由Nvidia做出部分貢獻的Python机器学习框架，用於变换数值函数。JAX结合了修改版的Autograd自動微分系統，以及來自OpenXLA專案的編譯器，可加速數值線性運算。其設計目標是在介面與程式設計風格上盡可能與NumPy保持相容，使使用者能夠以熟悉的方式撰寫高效能運算程式。此外，JAX亦可與TensorFlow、PyTorch等機器學習框架整合使用。

主要功能
JAX的主要功能是。

• Equinox，将参数化函数（包括人工神经网络）表示为PyTree的库。它由Patrick Kidger创建。
• Diffrax，用于求微分方程的数值解的库，比如解常微分方程和随机微分方程。
• Optax，DeepMind开发的用于梯度处理和最优化的库。
• Lineax，用于解线性方程组和。
• RLax，DeepMind开发的用于强化学习的库
• jraph，DeepMind开发的图神经网络库。
• jaxtyping，用于为阵列或张量的形状和数据类型增加类型标注的库。

*NumPyro，概率编程库。
*Brax，物理引擎。

参见

• NumPy
• TensorFlow
• PyTorch
• CUDA
• 自动微分
• 即时编译

*
*

引用
外部链接

• Documentationː
• Colab (Jupyter/iPython) Quickstart Guideː
• TensorFlow's XLAː (Accelerated Linear Algebra)

*

• Original paperː