jax.numpy.permute_dims

jax.numpy.permute_dims(a, /, axes)

置换数组的轴/维度。

JAX 对 array_api.permute_dims() 的实现。

参数:

• a (ArrayLike) – 输入数组

• axes (tuple[int, ...]) – 一个整数元组，范围为 [0, a.ndim)，指定轴的排列。

返回:

一个复制品 a，其中轴已被置换。

返回类型:

Array

另请参阅

• jax.numpy.transpose()

• jax.numpy.matrix_transpose()

示例

>>> a = jnp.array([[1, 2, 3],
...                [4, 5, 6]])
>>> jnp.permute_dims(a, (1, 0))
Array([[1, 4],
       [2, 5],
       [3, 6]], dtype=int32)