jax.vmap与jax.numpy.vectorize的区别、选择方式及性能差异

• 实现逻辑本质不同：jax.vmap是自动批处理原生原语，直接把函数运算逻辑映射到输入数组的批维度上，全程在JAX的tracing阶段完成，生成原生支持批运算的JAX中间表示。而jax.numpy.vectorize是对标NumPyvectorize接口的语法糖，本质是对输入数组的元素循环调用目标函数后拼接结果，只是循环被放到JAX runtime层面优化，并非原生批处理逻辑。
• 维度控制能力不同：jax.vmap支持通过in_axes、out_axes参数灵活指定输入输出的批处理维度，支持多输入多输出的复杂映射规则，还可嵌套实现多维度批处理。jax.numpy.vectorize默认仅支持逐元素映射，要指定批维度需额外配置signature参数，灵活性远低于vmap。
• 生态兼容性不同：jax.vmap作为JAX核心原语，可无缝和jax.jit/jax.grad/jax.pmap等其他变换组合，几乎没有兼容问题。jax.numpy.vectorize由于外层包裹了隐式循环，和其他JAX变换组合时的适配开销更高，部分复杂场景下会出现兼容问题。

• 仅做简单逐元素运算、需要对齐NumPy使用习惯、无复杂批处理需求时，可选择jax.numpy.vectorize，写法对NumPy用户更友好。
• 需要灵活控制批维度、要和其他JAX变换组合、后续有性能优化需求时，优先选jax.vmap，是JAX生态的标准批处理方案。
• 如果目标函数本身已经支持广播运算，两种都不需要用，直接传入数组即可，性能比两种向量化方案都高。

绝大多数场景下jax.vmap性能远优于jax.numpy.vectorize：
vmap生成的是原生批处理逻辑，编译后可以直接利用GPU/TPU等加速器的并行运算能力，没有额外循环开销。而jax.numpy.vectorize的隐式元素循环哪怕经过JIT优化，也会给编译器增加额外的优化负担，大数组、高维度场景下和vmap的性能差距会非常明显。

内容的提问来源于stack exchange，提问作者Jean-Eric