jax.numpy.matrix_transpose

jax.numpy.matrix_transpose(x, /)

转置数组的最后两个维度。

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

参数:

x (ArrayLike) – 输入数组，必须有 x.ndim >= 2

返回:

数组的矩阵转置副本。

返回类型:

Array

另请参阅

• jax.Array.mT: 通过 Array() 属性访问的相同操作。

• jax.numpy.transpose(): 一般的多轴转置

注意

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

示例

这是一个表示批处理 2x2 矩阵的 2x2x2 矩阵

>>> x = jnp.array([[[1, 2],
...                 [3, 4]],
...                [[5, 6],
...                 [7, 8]]])
>>> jnp.matrix_transpose(x)
Array([[[1, 3],
        [2, 4]],

       [[5, 7],
        [6, 8]]], dtype=int32)

为方便起见，您可以通过 mT 属性执行相同的转置 jax.Array

>>> x.mT
Array([[[1, 3],
        [2, 4]],

       [[5, 7],
        [6, 8]]], dtype=int32)