Google JAX

# 导入库from jax import gradimport 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)

# 导入库from jax import jitimport 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)

# 导入库from functools import partialfrom jax import vmapimport 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导入pmap和random；导入JAX NumPyfrom jax import pmap, randomimport jax.numpy as jnp# 生成2个维度为5000 x 6000的随机数矩阵，每设备一个random_keys = random.split(random.PRNGKey(0), 2)matrices = pmap(lambda key: random.normal(key, (5000, 6000)))(random_keys)# 没有数据传输，并行的在每个CPU/GPU上进行局部矩阵乘法 outputs = pmap(lambda x: jnp.dot(x, x.T))(matrices)# 没有数据传输，并行的在每个CPU/GPU上分别求取这两个矩阵的均值means = pmap(jnp.mean)(outputs)print(means)