jax_transpose

probly.representation.jax_functions.jax_transpose(a: JaxArrayLike, /, axes: Sequence[int] | None = None) → jax.Array

Return a transposed version of the array, mirroring jax.numpy.transpose.

Parameters:

• a – The array to transpose.

• axes – The permutation of axes. If omitted, the axes are reversed.

Returns:

The transposed array.