JAX介绍

        谷歌于 2018 年推出的 JAX 迎来了迅猛发展，2020 年，DeepMind 宣布使用 JAX 来加速其研究。越来越多来自谷歌大脑（Google Brain）和其他机构的项目也都在使用 JAX。 JAX 已经在深度学习、机器人 / 控制系统、贝叶斯方法和科学模拟等诸多领域得到了广泛应用。

        ​JAX 的定位科学计算（Scientific Computing）和函数转换（Function Transformations）的交叉融合，具有除训练深度学习模型以外的一系列能力，包括如下：

• 即时编译（Just-in-Time Compilation）
• 自动并行化（Automatic Parallelization）
• 自动向量化（Automatic Vectorization）
• 自动微分（Automatic Differentiation）

六条可能想要使用 JAX 的理由：

• NumPy 加速器。
• XLA。XLA（Accelerated Linear Algebra）是专为线性代数设计的全程序优化编译器。JAX 建立在 XLA 之上，显著提高了计算速度上限；
• JAX 允许用户使用 XLA 将自己的函数转换为即时编译（JIT）版本。这意味着可以通过在计算函数中添加一个简单的函数装饰器（decorator）来将计算速度提高几个数量级；
• Auto-differentiation。JAX 将 Autograd（自动区分原生 Python 代码和 NumPy 代码）和 XLA 结合在一起，它的自动微分能力在科学计算的许