如何在 jit 编译的 jax 代码中执行非整数索引运算？

【问题标题】：How to perform non-integer index arithmetic in jit-compiled jax code?如何在 jit 编译的 jax 代码中执行非整数索引运算？
【发布时间】：2021-03-16 19:15:54
【问题描述】：

如果我们对数组索引执行非整数计算（然后转换为 int() ），似乎我们仍然无法将结果用作 jit 编译的 jax 代码中的有效索引。我们如何解决这个问题？

以下是一个最小的示例。具体问题：命令 jnp.diag_indices(d) 是否可以在不向 fun() 传递额外参数的情况下工作

在木星单元中运行：

import jax.numpy as jnp
from jax import jit

@jit
def fun(t):
    d = jnp.sqrt(t.size**2)
    d = jnp.array(d,int)
    
    jnp.diag_indices(t.size)   # this line works
    jnp.diag_indices(d)        # this line breaks. Comment it out to see that d and t.size have the same dtype=int32 

    return t.size, d
    
fun(jnp.array([1,2]))

【问题讨论】：

标签： python jax

【解决方案1】：

问题不在于d 的类型，而是d 是jax 操作的结果，因此在JIT 上下文中跟踪。在 JAX 中，数组的形状和大小不能依赖于跟踪的数量，这就是您的代码会导致错误的原因。

要解决这个问题，一个有用的模式是使用np 操作而不是jnp 操作来确保d 是静态的且不被跟踪：

import jax.numpy as jnp
from jax import jit

@jit
def fun(t):
    d = np.sqrt(t.size**2)
    d = np.array(d, int)
    
    jnp.diag_indices(t.size)
    jnp.diag_indices(d)

    return t.size, d
    
print(fun(jnp.array([1,2])))
# (DeviceArray(2, dtype=int32), DeviceArray(2, dtype=int32))

有关跟踪、静态值和类似主题的简要背景信息，How To Think In JAX 文档页面可能会有所帮助。

【讨论】：