jax_swapaxes

probly.representation.jax_functions.jax_swapaxes(a: JaxArrayLike, /, axis1: int, axis2: int) → jax.Array

Swap two axes of the array, mirroring jax.numpy.swapaxes.

Parameters:

• a – The array to swap axes of.

• axis1 – The first axis.

• axis2 – The second axis.

Returns:

The array with the two axes swapped.