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
