JAX 中高效的自定义数组创建例程

Question

我仍在

jax

中掌握最佳实践。我的广泛问题如下：

在

jax

中实现自定义数组创建例程的最佳实践是什么？

例如，我想实现一个函数，该函数创建一个矩阵，除了给定列中的 1 之外，到处都是 0。我选择了这个（Jupyter 笔记本）：

import numpy as np
import jax.numpy as jnp

def ones_at_col(shape_mat, idx):
    idxs = jnp.arange(shape_mat[1])[None,:]
    mat = jnp.where(idx==idxs, 1, 0)
    mat = jnp.repeat(mat, shape_mat[0], axis=0)
    return mat

shape_mat = (5,10)

print(ones_at_col(shape_mat, 5))

%timeit np.zeros(shape_mat)

%timeit jnp.zeros(shape_mat)

%timeit ones_at_col(shape_mat, 5)

输出是

[[0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]
 [0 0 0 0 0 1 0 0 0 0]]
127 ns ± 0.717 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)
31.3 µs ± 331 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
123 µs ± 1.79 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

我的函数比

jnp.zeros()

例程慢 4 倍，这还不错。这告诉我我所做的并不疯狂。

但是两个

jax

例程都比同等的

numpy

例程慢得多。这些函数无法进行 jitted，因为它们将形状作为参数，因此无法跟踪。我想这就是为什么它们本质上较慢的原因？我猜想，如果它们中的任何一个出现在另一个 jitted 函数的范围内，它们就可以被追踪并加速？

我可以做更好的事情吗？还是我正在挑战

jax

的极限？

Answer 1

最好的方法可能是这样的：

mat = jnp.zeros(shape_mat).at[:, 5].set(1)

关于与 NumPy 的时序比较，相关阅读是 JAX FAQ：JAX 比 NumPy 更快吗？ 总结是，对于这种特殊情况（创建一个简单的数组），您不会期望 JAX 在性能方面与 NumPy 匹配。

如果您希望在 JAX 中获得更快的性能，您应该做的另一件事是使用

jax.jit

。例如，这个版本的函数应该是相当最优的（尽管同样，由于常见问题解答链接中讨论的原因，它的速度不如 NumPy）：

@partial(jax.jit, static_argnames=['shape_mat', 'idx'])
def ones_at_col(shape_mat, idx):
  return jnp.zeros(shape_mat).at[:, idx].set(1)

如果您要使用不同的索引值多次调用此函数，则可以将

idx

设置为非静态。

另一个旁注：您的微基准测试可能无法测量您认为它们正在测量的内容：有关这方面的提示，请参阅JAX 常见问题解答：对 JAX 代码进行基准测试。特别要注意编译时间和异步调度效果。