Skip to content

take

Returns elements of an array along an axis.

Parameters:

Name Type Description Default
x SparseArray

Input array.

required
indices ndarray

Array indices. The array must be one-dimensional and have an integer data type.

required
axis int

Axis over which to select values. If axis is negative, the function must determine the axis along which to select values by counting from the last dimension. For None, the flattened input array is used. Default: None.

None

Returns:

Name Type Description
out COO

A COO array with requested indices.

Raises:

Type Description
ValueError

If the input array isn't and can't be converted to COO format.

Source code in sparse/numba_backend/_coo/common.py
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
def take(x, indices, /, *, axis=None):
    """
    Returns elements of an array along an axis.

    Parameters
    ----------
    x : SparseArray
        Input array.
    indices : ndarray
        Array indices. The array must be one-dimensional and have an integer data type.
    axis : int
        Axis over which to select values. If ``axis`` is negative, the function must
        determine the axis along which to select values by counting from the last dimension.
        For ``None``, the flattened input array is used. Default: ``None``.

    Returns
    -------
    out : COO
        A COO array with requested indices.

    Raises
    ------
    ValueError
        If the input array isn't and can't be converted to COO format.
    """

    x = _validate_coo_input(x)

    if axis is None:
        x = x.flatten()
        return x[indices]

    axis = normalize_axis(axis, x.ndim)
    full_index = (slice(None),) * axis + (indices, ...)
    return x[full_index]