用于简单数组更新的Jax-vmap

Question

我是Jax的新手，我正在转换其他人的代码，该代码使用了numba"fastmath"；特性，并依赖于许多嵌套的for循环，而没有太多性能损失。我正在尝试使用Jax的vmap函数重新创建相同的行为。然而，我目前正为一些基本问题而苦苦挣扎。下面是一个我试图使用vmap进行矢量化的简化示例：

import jax.numpy as jnp
from jax import vmap
import jax.ops
a = jnp.arange(20).reshape((4, 5))
b = jnp.arange(5)
c = jnp.arange(4)
d = jnp.zeros(20)
e = jnp.zeros((4, 5))
for i in range(a.shape[0]):
for j in range(a.shape[1]):
a = jax.ops.index_add(a, jax.ops.index[i, j], b[j] + c[i])
d = jax.ops.index_update(d, jax.ops.index[i*a.shape[1] + j], b[j] * c[i])
e = jax.ops.index_update(e, jax.ops.index[i, j], 2*b[j])

如何使用vmap重写这样的代码？虽然这段代码相对容易手动向量化，但我希望更好地了解vmap是如何工作的，并希望任何答案都能对我有所帮助。文档现在似乎并没有真正帮助我。我真的很感激你能提供的任何帮助。

Answer 1

以下是如何使用vmap:实现大致相同的计算

from jax import vmap, partial
@partial(vmap, in_axes=(0, None, 0))
@partial(vmap, in_axes=(0, 0, None))
def f(a, b, c):
return a + b + c, b * c, 2 * b
a, d, e = f(a, b, c)
d = d.ravel()