fn

bucketize

Tensor
bucketize(values: Tensor, boundaries: Tensor, right: bool = ...)
source

Return the bucket index of each element of values in boundaries.

For each v = values[i] returns the smallest j such that v < boundaries[j] (or v <= boundaries[j] when right=True). Equivalent to a per-element binary search.

Parameters

valuesTensor
Query values (arbitrary shape).
boundariesTensor
1-D, sorted, monotonically non-decreasing bucket edges.
rightbool= False
Use <= rather than < for the comparison.

Returns

Tensor

int64 tensor of bucket indices with the same shape as values.

Notes

Indices are in [0, len(boundaries)] — out-of-range queries map to the boundary slots. Sister function searchsorted has the same semantics but takes the sorted sequence as the first argument.

Examples

>>> import lucid
>>> bounds = lucid.tensor([1.0, 3.0, 5.0])
>>> lucid.bucketize(lucid.tensor([0.5, 2.0, 4.0, 6.0]), bounds)
Tensor([0, 1, 2, 3])