从函数有效地填充数组

Question

我想从一个函数构造一个二维数组，这样我就可以利用 jax.jit。

我通常使用 numpy 执行此操作的方法是创建一个空数组，然后就地填充该数组。

xx = jnp.empty((num_a, num_b))
yy = jnp.empty((num_a, num_b))
zz = jnp.empty((num_a, num_b))

for ii_a in range(num_a):
    for ii_b in range(num_b):
        a = aa[ii_a, ii_b]
        b = bb[ii_a, ii_b]

        xyz = self.get_coord(a, b)

        xx[ii_a, ii_b] = xyz[0]
        yy[ii_a, ii_b] = xyz[1]
        zz[ii_a, ii_b] = xyz[2]

为了在 jax 内完成这项工作，我尝试使用 jax.opt.index_update。

        xx = xx.at[ii_a, ii_b].set(xyz[0])
        yy = yy.at[ii_a, ii_b].set(xyz[1])
        zz = zz.at[ii_a, ii_b].set(xyz[2])

它运行没有错误，但是当我尝试使用 @jax.jit 装饰器时速度非常慢（至少比纯 python/numpy 版本慢一个数量级）。

使用 jax 从函数填充多维数组的最佳方法是什么？

Answer 1

JAX 有一个 vmap transform 是专门为这种应用程序设计的。

只要您的 get_coords 函数与 JAX 兼容（即是一个没有 side-effects 的纯函数），您可以在一行中完成此操作：

from jax import vmap
xx, yy, zz = vmap(vmap(get_coord))(aa, bb)

Answer 2

这可以通过使用 jax.vmap or the jax.numpy.vectorize 函数来有效地实现。

使用vectorize的例子：

import jax.numpy as jnp

def get_coord(a, b):
    return jnp.array([a, b, a+b])

f0 = jnp.vectorize(get_coord, signature='(),()->(i)')
f1 = jnp.vectorize(f0, excluded=(1,), signature='()->(i,j)')

xyz = f1(a,b)

vectorize 函数在底层使用了 vmap，所以这应该完全等同于：

f0 = jax.vmap(get_coord, (None, 0))
f1 = jax.vmap(f0, (0, None))

使用vectorize的好处是代码仍然可以是标准numpy中的运行。缺点是代码不太简洁，并且由于包装器可能会有少量开销。

从函数有效地填充数组

Efficiently fill an array from a function

python

numpy

jax