python - 如何将 jax vmap 用于嵌套循环？

Question

我想使用 vmap 来矢量化此代码以提高性能。

def matrix(dataA, dataB):
    return jnp.array([[func(a, b) for b in dataB] for a in dataA])
matrix(data, data)

我试过这个:

def f(x, y):
    return func(x, y)
mapped = jax.vmap(f)
mapped(data, data)

但这仅给出对角线条目。
基本上我有一个向量 data = [1,2,3,4,5] (示例)，我想得到一个矩阵，使得每个条目 (i, j)矩阵是 f(data[i], data[j]) .因此，生成的矩阵形状将为 (len(data), len(data)) .

Answer 1

jax.vmap一次映射一组轴。如果要映射两个独立的轴集，可以通过嵌套两个 vmap 来实现。转换:

mapped = jax.vmap(jax.vmap(f, in_axes=(None, 0)), in_axes=(0, None))
result = mapped(data, data)