JAX的深度学习和科学计算

         JAX 框架：能使用GPU加速的，支持自动微分的numpy。numpy在科学计算领域十分普及，但在深度学习领域。

         Jax 是 Autograd 和 XLA 的结合用以提供高性能机器学习研究。作为 Autograd 的更新版本，JAX 可以自动微分原生 Python 和 Numpy 函数。可通过 loops、branches、recursion 和 closures，进行微分，并且能够对导数的导数求导。

        2020 年，DeepMind 宣布使用 JAX 来加速其研究。越来越多来自谷歌大脑（Google Brain）和其他机构的项目也都在使用 JAX

        JAX 是一个非常有前途的项目，并且用户一直在稳步增长。JAX 已经在深度学习、机器人 / 控制系统、贝叶斯方法和科学模拟等诸多领域得到了广泛应用

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

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

JAX 的速度比 NumPy 快了 N 个数量级，JAX 使用的是 TPU，NumPy 使用了 CPU，以此强调 JAX 的速度上限远高于 NumPy

JAX 包含了一个可扩展系统来实现这样的函数转换，有四种典型方式：

        Grad() 进行自动微分；

        Vmap() 自动向量化；

        Pmap() 并行化计算；

        Jit() 将函数转换为即时编译版本

科学计算

 深度学习

深度学习入门

参考：

Why You Should (or Shouldn't) be Using Google's JAX in 2022 (assemblyai.com)