1

以下 numpy 代码非常好：

arr = np.arange(50)
print(arr.shape) # (50,)

indices = np.zeros((30,), dtype=int)
print(indices.shape) # (30,)

arr[indices]

迁移到 jax 后它也可以工作：

import jax.numpy as jnp

arr = jnp.arange(50)
print(arr.shape) # (50,)

indices = jnp.zeros((30,), dtype=int)
print(indices.shape) # (30,)

arr[indices]

现在让我们尝试混合使用 numpy 和 jax：

arr = np.arange(50)
print(arr.shape) # (50,)

indices = jnp.zeros((30,), dtype=int)
print(indices.shape) # (30,)

arr[indices]

这会产生以下错误：

IndexError: too many indices for array: array is 1-dimensional, but 30 were indexed

如果不支持使用 jax 数组索引到 numpy 数组，那对我来说很好。但是错误信息似乎是错误的。事情变得更加混乱。如果稍微改变形状，代码就可以正常工作。在下面的示例中，我只编辑了从 (30,) 到 (40,) 的索引形状。没有更多错误消息：

arr = np.arange(50)
print(arr.shape) # (50,)

indices = jnp.zeros((40,), dtype=int)
print(indices.shape) # (40,)

arr[indices]

我在 cpu 上运行 jax 版本“0.2.12”。这里发生了什么？

4

1 回答 1

1

这是一个长期存在的已知问题（请参阅https://github.com/google/jax/issues/620）；这不是 JAX 本身可以轻松修复的错误，但需要更改 NumPy 处理非ndarray索引的方式。好消息是修复即将到来：上面有问题的代码伴随着以下警告，该警告来自 NumPy：

FutureWarning: Using a non-tuple sequence for multidimensional indexing is
 deprecated; use `arr[tuple(seq)]` instead of `arr[seq]`. In the future this
 will be interpreted as an array index, `arr[np.array(seq)]`, which will result
 either in an error or a different result.

一旦这个弃用周期完成，JAX 数组将在 NumPy 索引中正常工作。

np.asarray在此之前，您可以通过在使用 JAX 数组索引 NumPy 数组时显式调用来解决它。

于 2021-05-30T03:42:55.887 回答