Function permute_dims

xarray/namedarray/_array_api.py:195–221

Permutes the dimensions of an array. Parameters ---------- x : Array to permute. axes : Permutation of the dimensions of x. Returns ------- out : An array with permuted dimensions. The returned array must have the same data type as x

(x: NamedArray[Any, _DType], axes: _Axes)

def permute_dims(x: NamedArray[Any, _DType], axes: _Axes) -> NamedArray[Any, _DType]:
    """
    Permutes the dimensions of an array.

    Parameters
    ----------
    x :
        Array to permute.
    axes :
        Permutation of the dimensions of x.

    Returns
    -------
    out :
        An array with permuted dimensions. The returned array must have the same
        data type as x.

    """

    dims = x.dims
    new_dims = tuple(dims[i] for i in axes)
    if isinstance(x._data, _arrayapi):
        xp = _get_data_namespace(x)
        out = x._new(dims=new_dims, data=xp.permute_dims(x._data, axes))
    else:
        out = x._new(dims=new_dims, data=x._data.transpose(axes))  # type: ignore[attr-defined]
    return out