jax.numpy.transpose

jax.numpy.transpose(a, axes=None)

返回 N 维数组的转置版本。

JAX 实现的 numpy.transpose()，使用 jax.lax.transpose() 实现。

参数:

• a (ArrayLike) – 输入数组

• axes (Sequence[int] | None | None) – 可选地指定使用长度为 a.ndim 的整数序列 i 的排列，满足 0 <= i < a.ndim。默认为 range(a.ndim)[::-1]，即反转所有轴的顺序。

返回:

数组的转置副本。

返回类型:

数组

参见

• jax.Array.transpose()：通过 Array 方法的等效函数。

• jax.Array.T：通过 Array 属性的等效函数。

• jax.numpy.matrix_transpose()：转置数组的最后两个轴。这适用于处理批量 2D 矩阵。

• jax.numpy.swapaxes()：交换数组中的任意两个轴。

• jax.numpy.moveaxis()：将轴移动到数组中的另一个位置。

注意

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

示例

对于 1D 数组，转置是恒等变换

>>> x = jnp.array([1, 2, 3, 4])
>>> jnp.transpose(x)
Array([1, 2, 3, 4], dtype=int32)

对于 2D 数组，转置是矩阵转置

>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> jnp.transpose(x)
Array([[1, 3],
       [2, 4]], dtype=int32)

对于 N 维数组，转置反转轴的顺序

>>> x = jnp.zeros(shape=(3, 4, 5))
>>> jnp.transpose(x).shape
(5, 4, 3)

可以指定 axes 参数来更改此默认行为

>>> jnp.transpose(x, (0, 2, 1)).shape
(3, 5, 4)

由于交换最后两个轴是常见的操作，因此可以通过其自身的 API jax.numpy.matrix_transpose() 完成

>>> jnp.matrix_transpose(x).shape
(3, 5, 4)

为方便起见，也可以使用 jax.Array.transpose() 方法或 jax.Array.T 属性执行转置

>>> x = jnp.array([[1, 2],
...                [3, 4]])
>>> x.transpose()
Array([[1, 3],
       [2, 4]], dtype=int32)
>>> x.T
Array([[1, 3],
       [2, 4]], dtype=int32)