jax_moveaxis¶

probly.representation.jax_functions.jax_moveaxis(a: JaxArrayLike, /, source: int | Sequence[int], destination: int | Sequence[int]) → jax.Array[source]¶

Move axes of the array to new positions, mirroring jax.numpy.moveaxis.

Parameters:

a – The array to move axes of.
source – The axes to move.
destination – The destination positions of the moved axes.

Returns:

The array with the axes moved.