JAX

JAX，是用于变换数值函数的Python机器学习框架，它由Google开发并具有来自Nvidia的一些贡献。它结合了修改版本的Autograd（自动通过函数的微分获得其梯度函数），和OpenXLA的（加速线性代数）。它被设计为尽可能的遵从NumPy的结构和工作流程，并协同工作于各种现存的框架如TensorFlow和PyTorch。

主要功能
JAX的主要功能是：
grad：自动微分，
jit：即时编译，
vmap：自动向量化，
pmap：编程。

 grad 

下面的代码演示函数的自动微分。 

# 导入库
from jax import grad
import jax.numpy as jnp

# 定义logistic函数
def logistic(x):  
    return jnp.exp(x) / (jnp.exp(x) + 1)

# 获得logistic函数的梯度函数
grad_logistic = grad(logistic)

# 求值logistic函数在x = 1处的梯度 
grad_log_out = grad_logistic(1.0)   
print(grad_log_out) 

最终的输出为：
0.19661194

 jit 

下面的代码演示jit函数的优化。

# 导入库
from jax import jit
import jax.numpy as jnp

# 定义cube函数
def cube(x):
    return x * x * x

# 生成数据
x = jnp.ones((10000, 10000))

# 创建cube函数的jit版本
jit_cube = jit(cube)

# 应用cube函数和jit_cube函数于相同数据来比较其速度
cube(x)
jit_cube(x)

可见的运行时间显著的短于。

 vmap 

下面的代码展示函数的通过SIMD的向量化。

# 导入库
from functools import partial
from jax import vmap
import jax.numpy as jnp

# 定义函数
def grads(self, inputs):
    in_grad_partial = partial(self._net_grads, self._net_params)
    grad_vmap = vmap(in_grad_partial)
    rich_grads = grad_vmap(inputs)
    flat_grads = np.asarray(self._flatten_batch(rich_grads))
    assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
    return flat_grads 

 使用JAX的库 
一些Python库使用JAX作为后端，这包括：

 Flax，最初由Google Brain开发的高层人工神经网络库。
 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ː 

Category:数据挖掘和机器学习软件
Category:机器学习
Category:Google軟體
Category:用Python編程的自由軟體
Category:Python科学库
Category:使用Apache许可证的软件
Category:开源人工智能