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
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
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]