JAX深度学习框架

GoogleJAX是一个专为数值计算优化的机器学习框架，Google将其描述为融合了改进版Autograd（通过自动微分功能生成梯度）与TensorFlow的XLA（加速线性代数计算）。这一框架的设计理念尽量模仿NumPy的结构和操作流程，并且能够与TensorFlow、PyTorch等多种现有框架无缝兼容。

JAX的核心功能包括：

• grad：提供自动微分功能
• jit：支持即时编译
• vmap：实现自动矢量化
• pmap：支持SPMD（单指令多数据）编程模式