expand_dims

扩展数组的形状，通过在由 axis 指定的位置插入一个大小为一的新轴（维度）。

参数

名称	类型	描述	默认值
a	COO	输入的 COO 数组。	必需
axis	int	新轴在扩展后的轴中的放置位置。	0

返回值

名称	类型	描述
result	COO	一个具有与 x 相同数据类型的扩展后的输出 COO 数组。

示例

>>> 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)

源代码位于 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,
    )