### 介绍 JAX 中的 `vmap` 函数 JAX 提供了多种函数变换工具，其中 `vmap` 是一种用于向量化操作的核心机制。它通过将函数自动映射到输入数组的批量维度上，从而实现对单个样本处理函数的批量处理能力。这种技术可以显著简化代码，并提高计算效率，尤其是在处理大量数据时。 `vmap` 的基本形式为： ```python jax.vmap(fun, in_axes=0, out_axes=0) ``` - `fun`: 要向量化的函数。 - `in_axes`: 指定输入参数中哪些轴将被向量化（即批量处理）。 - `out_axes`: 指定输出结果中哪些轴对应于批量维度。 例如，假设有一个仅能处理单个一维向量的函数，使用 `vmap` 可以将其扩展为支持二维矩阵输入，其中每一行被视为一个独立样本[^1]。 ### 示例用法 以下是一个具体的示例，展示如何使用 `vmap` 将一个只能处理单个样本的函数扩展为支持批量处理： ```python import jax import jax.numpy as jnp # 定义一个只能处理单个样本的函数 def sample_function(x): return jnp.sum(x ** 2) # 使用 vmap 扩展该函数以支持批量处理 batched_function = jax.vmap(sample_function, in_axes=0) # 输入数据：批量大小为 4，每个样本是 3 维向量 data = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0], [10.0, 11.0, 12.0]]) # 调用批量处理函数 result = batched_function(data) print("Batched result:", result) ``` 运行上述代码后，输出结果为： ``` Batched result: [ 14. 77. 194. 365.] ``` 这表明 `vmap` 成功地将原本只能处理单个样本的 `sample_function` 扩展为能够处理四组输入数据的能力，且每组数据的结果独立计算并堆叠在一起[^1]。 ### 高阶用法：多输入与自定义轴 `vmap` 也支持多输入函数，并允许指定不同的 `in_axes` 来控制不同参数的向量化方式。例如： ```python def multi_input_function(x, y): return x * y + x # 对第一个输入的第 0 轴进行向量化，第二个输入保持不变 batched_multi_input = jax.vmap(multi_input_function, in_axes=(0, None)) x_data = jnp.array([1.0, 2.0, 3.0]) y_value = jnp.array(5.0) result = batched_multi_input(x_data, y_value) print("Result with mixed axes:", result) ``` 输出结果为： ``` Result with mixed axes: [ 6. 12. 18.] ``` 这里，`x_data` 被视为批量数据，而 `y_value` 被广播到所有批次中，实现了灵活的向量化行为[^1]。 ### 性能优势与应用场景 在实际应用中，`vmap` 不仅简化了代码逻辑，还能够利用底层 XLA 编译优化，显著提升计算性能。相比手动编写循环或使用 NumPy 进行批量处理，JAX 的 `vmap` 更加高效，尤其适用于大规模数据集上的模型推理和训练任务[^4]。 此外，`vmap` 与 `grad`、`jit` 等其他 JAX 工具结合使用时效果更佳。例如，在批量计算梯度时，可以通过嵌套 `vmap` 和 `grad` 实现高效的 Hessian 计算或批量自动微分需求[^2]。 ---