JAX（一）

> JAX 是一个用于高性能数值计算的 Python 库，特别为机器学习领域的高性能计算设计。它的 API 基于 Numpy 构建，包含丰富的数值计算与科学计算函数。JAX其实是 TensorFlow 的一个简化库，结合 Autograd 和 XLA，可以支持部分 TensorFlow 的功能，但是比 TensorFlow 更加简洁易用。

> Python 和 Numpy 的广泛使用，使得 JAX 十分简洁、灵活、易于上手，学习成本也比较低。除了 Numpy 的 API 外，JAX 还包含一系列可拓展、可组合的系统功能，有力地支持了机器学习研究。这些功能特性主要包括：

- 可差分：基于梯度的优化方法在机器学习领域具有十分重要的作用。JAX 可通过grad、hessian、jacfwd 和 jacrev 等函数转换，原生支持任意数值函数的前向和反向模式的自动微分。

- 向量化：在机器学习中，通常需要在大规模的数据上运行相同的函数，例如计算整个批次的损失或每个样本的损失等。JAX 通过 vmap 变换提供了自动矢量化算法，大大简化了这种类型的计算，这使得研究人员在处理新算法时无需再去处理批量化的问题。JAX 同时还可以通过 pmap 转换支持大规模的数据并行，从而优雅地将单个处理器无法处理的大数据进行处理。

- JIT编译：XLA (Accelerated Linear Algebra, 加速线性代数) 被用于 JIT 即时编译，在 GPU 和云 TPU 加速器上执行 JAX 程序。JIT 编译与 JAX 的 API (与 Numpy 一致的数据函数) 为研发人员提供了便捷接入高性能计算的可能，无需特别的经验就能将计算运行在多个加速器上。