torch.searchsorted

torch.searchsorted(sorted_sequence, values, *, out_int32=False, right=False, side='left', out=None, sorter=None) → Tensor

Find the indices from the innermost dimension of sorted_sequence such that, if the corresponding values in values were inserted before the indices, when sorted, the order of the corresponding innermost dimension within sorted_sequence would be preserved. Return a new tensor with the same size as values. If right is False or side is ‘left (default), then the left boundary of sorted_sequence is closed. More formally, the returned index satisfies the following rules:

sorted_sequence	right	returned index satisfies
1-D	False	sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]
1-D	True	sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]
N-D	False	sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]
N-D	True	sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]

Parameters

• sorted_sequence (Tensor) – N-D or 1-D tensor, containing monotonically increasing sequence on the innermost dimension unless sorter is provided, in which case the sequence does not need to be sorted

• values (Tensor or Scalar) – N-D tensor or a Scalar containing the search value(s).

Keyword Arguments

• out_int32 (bool, optional) – indicate the output data type. torch.int32 if True, torch.int64 otherwise. Default value is False, i.e. default output data type is torch.int64.

• right (bool, optional) – if False, return the first suitable location that is found. If True, return the last such index. If no suitable index found, return 0 for non-numerical value (eg. nan, inf) or the size of innermost dimension within sorted_sequence (one pass the last index of the innermost dimension). In other words, if False, gets the lower bound index for each value in values on the corresponding innermost dimension of the sorted_sequence. If True, gets the upper bound index instead. Default value is False. side does the same and is preferred. It will error if side is set to “left” while this is True.

• side (str, optional) – the same as right but preferred. “left” corresponds to False for right and “right” corresponds to True for right. It will error if this is set to “left” while right is True.

• out (Tensor, optional) – the output tensor, must be the same size as values if provided.

• sorter (LongTensor, optional) – if provided, a tensor matching the shape of the unsorted sorted_sequence containing a sequence of indices that sort it in the ascending order on the innermost dimension

Example:

>>> sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]])
>>> sorted_sequence
tensor([[ 1,  3,  5,  7,  9],
        [ 2,  4,  6,  8, 10]])
>>> values = torch.tensor([[3, 6, 9], [3, 6, 9]])
>>> values
tensor([[3, 6, 9],
        [3, 6, 9]])
>>> torch.searchsorted(sorted_sequence, values)
tensor([[1, 3, 4],
        [1, 2, 4]])
>>> torch.searchsorted(sorted_sequence, values, side='right')
tensor([[2, 3, 5],
        [1, 3, 4]])

>>> sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9])
>>> sorted_sequence_1d
tensor([1, 3, 5, 7, 9])
>>> torch.searchsorted(sorted_sequence_1d, values)
tensor([[1, 3, 4],
        [1, 3, 4]])