swapaxes

class brainunit.math.swapaxes(a, axis1, axis2, **kwargs)

Interchange two axes of a quantity or an array.

Parameters:

• a (Array | saiunit.Quantity) – Input array.

• axis1 (int) – First axis.

• axis2 (int) – Second axis.

Returns:

a_swapped – a new array where the axes are swapped.

Return type:

Array | saiunit.Quantity

Examples

>>> import saiunit as u
>>> import jax.numpy as jnp
>>> a = jnp.zeros((3, 4, 5)) * u.meter
>>> u.math.swapaxes(a, 0, 2).shape
(5, 4, 3)