jax.vmap和jax.pmap介绍

jax.vmap 和 jax.pmap 是 JAX 中用于在不同层面进行并行计算的重要工具，它们有助于提高深度学习模型的效率。

import jax
import jax.numpy as jnp

# 定义一个简单的函数
def example_fn(x):
    return jnp.sum(x**2)

# 输入数据
key = jax.random.PRNGKey(42)
x_data = jax.random.uniform(key, shape=(10, 10))


### 1. jax.vmap
## jax.vmap 是向量化映射的缩写，它用于将一个函数向量化，以便能够同时处理批量的输入
# 使用 jax.vmap 对函数进行向量化
vectorized_fn = jax.vmap(example_fn)
result = vectorized_fn(x_data)
print(result.shape)  # (10,)
print(result) 

### 2. jax.pmap
## jax.pmap 是并行映射的缩写，它用于在多个设备上并行地对不同的输入执行相同的函数。
# 主要用于在多设备上进行模型并行计算。
parallel_fn = jax.pmap(example_fn) 

# print(jax.local_device_count()) 

# 应用并行映射函数
# 单机运行会报错
result = parallel_fn(x_data)
print(result)

参考：

https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html

https://jax.readthedocs.io/en/latest/_autosummary/jax.pmap.html