jax.numpy.rollaxis

jax.numpy.rollaxis(a, axis, start=0)

将指定的轴滚动到给定的位置。

numpy.rollaxis() 的 JAX 实现。

此函数为了与 NumPy 兼容而存在，但在大多数情况下，应使用更新的 jax.numpy.moveaxis() 代替，因为其参数的含义更直观。

参数:

• a (ArrayLike) – 输入数组。

• axis ( int ) – 要向前滚动的轴的索引。

• start ( int ) – 轴将滚动到的索引（默认 = 0）。在规范化负轴之后，如果 start <= axis，则轴滚动到 start 索引；如果 start > axis，则轴滚动到 start 之前的位置。

返回:

具有滚动轴的 a 的副本。

返回类型:

数组

注意事项

与 numpy.rollaxis() 不同，jax.numpy.rollaxis() 将返回输入数组的副本，而不是视图。但是，在 JIT 下，编译器将在可能的情况下优化掉这些副本，因此这在实践中不会对性能产生影响。

另请参阅

• jax.numpy.moveaxis()：具有比 rollaxis 更清晰语义的较新 API；在大多数情况下，应优先使用此函数而不是 rollaxis。

• jax.numpy.swapaxes()：交换两个轴。

• jax.numpy.transpose()：轴的一般排列。

示例

>>> a = jnp.ones((2, 3, 4, 5))

将轴 2 滚动到数组的开头

>>> jnp.rollaxis(a, 2).shape
(4, 2, 3, 5)

将轴 1 滚动到数组的末尾

>>> jnp.rollaxis(a, 1, a.ndim).shape
(2, 4, 5, 3)

以下两个与 moveaxis() 等效

>>> jnp.moveaxis(a, 2, 0).shape
(4, 2, 3, 5)
>>> jnp.moveaxis(a, 1, -1).shape
(2, 4, 5, 3)