jax.numpy.矩阵_转置

jax.numpy.矩阵_转置#

jax.numpy.matrix_transpose(x, /)[源代码][源代码]#

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

JAX 实现的 numpy.matrix_transpose(),基于 jax.lax.transpose() 实现。

参数:

x (ArrayLike) – 输入数组,必须满足 x.ndim >= 2

返回:

矩阵-数组的转置副本。

返回类型:

Array

参见

备注

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)

为了方便,你可以通过 jax.ArraymT 属性执行相同的转置:

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

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