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
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
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]