jax_matrix_transpose

probly.representation.jax_functions.jax_matrix_transpose(x: JaxArrayLike, /) → jax.Array

Transpose the last two axes of the array, mirroring jax.numpy.matrix_transpose.

Parameters:

x – The array to transpose.

Returns:

The array with its last two axes swapped.