1

我想以我可以利用的方式从函数构造一个二维数组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？

4

2 回答 2

1

JAX 具有专门为此类应用程序设计的vmap转换。

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

from jax import vmap
xx, yy, zz = vmap(vmap(get_coord))(aa, bb)
于 2020-10-07T17:20:42.497 回答
0

这可以通过使用jax.vmap或 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中运行。缺点是代码不太简洁，并且由于包装器可能会产生少量开销。

于 2020-10-15T21:08:24.070 回答