expand_dims

Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by axis.

Parameters:

Name	Type	Description	Default
a	COO	Input COO array.	required
axis	int	Position in the expanded axes where the new axis is placed.	0

Returns:

Name	Type	Description
result	COO	An expanded output COO array having the same data type as x.

Examples:

>>> import sparse
>>> x = sparse.COO.from_numpy([[1, 0, 0, 0, 2, -3]])
>>> x.shape
(1, 6)
>>> y1 = sparse.expand_dims(x, axis=1)
>>> y1.shape
(1, 1, 6)
>>> y2 = sparse.expand_dims(x, axis=2)
>>> y2.shape
(1, 6, 1)

Source code in sparse/numba_backend/_coo/common.py

def expand_dims(x, /, *, axis=0):
    """
    Expands the shape of an array by inserting a new axis (dimension) of size
    one at the position specified by ``axis``.

    Parameters
    ----------
    a : COO
        Input COO array.
    axis : int
        Position in the expanded axes where the new axis is placed.

    Returns
    -------
    result : COO
        An expanded output COO array having the same data type as ``x``.

    Examples
    --------
    >>> import sparse
    >>> x = sparse.COO.from_numpy([[1, 0, 0, 0, 2, -3]])
    >>> x.shape
    (1, 6)
    >>> y1 = sparse.expand_dims(x, axis=1)
    >>> y1.shape
    (1, 1, 6)
    >>> y2 = sparse.expand_dims(x, axis=2)
    >>> y2.shape
    (1, 6, 1)

    """

    x = _validate_coo_input(x)

    if not isinstance(axis, int):
        raise IndexError(f"Invalid axis position: {axis}")

    axis = normalize_axis(axis, x.ndim + 1)

    new_coords = np.insert(x.coords, obj=axis, values=np.zeros(x.nnz, dtype=np.intp), axis=0)
    new_shape = list(x.shape)
    new_shape.insert(axis, 1)
    new_shape = tuple(new_shape)

    from .core import COO

    return COO(
        new_coords,
        x.data,
        shape=new_shape,
        fill_value=x.fill_value,
    )