torch.searchsorted

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

從 sorted_sequence 的*最內層*維度中查詢索引，使得如果將 values 中相應的值插入到這些索引之前，排序後，sorted_sequence 中相應*最內層*維度的順序得以保持。返回一個與 values 大小相同的新張量。更正式地說，返回的索引滿足以下規則

sorted_sequence	right	返回的索引滿足
1D	False	sorted_sequence[i-1] < values[m][n]...[l][x] <= sorted_sequence[i]
1D	True	sorted_sequence[i-1] <= values[m][n]...[l][x] < sorted_sequence[i]
ND	False	sorted_sequence[m][n]...[l][i-1] < values[m][n]...[l][x] <= sorted_sequence[m][n]...[l][i]
ND	True	sorted_sequence[m][n]...[l][i-1] <= values[m][n]...[l][x] < sorted_sequence[m][n]...[l][i]

引數

• sorted_sequence (Tensor) – N 維或 1 維張量，在其*最內層*維度上包含單調遞增的序列，除非提供了 sorter，在這種情況下序列無需排序

• values (Tensor 或 Scalar) – N 維張量或標量，包含要搜尋的值。

關鍵字引數

• out_int32 (bool, optional) – 指示輸出資料型別。如果為 True，則為 torch.int32；否則為 torch.int64。預設值為 False，即預設輸出資料型別為 torch.int64。

• right (bool, optional) – 如果為 False，則返回找到的第一個合適位置。如果為 True，則返回最後一個此類索引。如果沒有找到合適的索引，對於非數值（例如 nan, inf）返回 0，或者返回 sorted_sequence 中*最內層*維度的大小（比*最內層*維度的最後一個索引大一）。換句話說，如果為 False，則獲取 values 中每個值在 sorted_sequence 相應*最內層*維度上的下界索引。如果為 True，則獲取上界索引。預設值為 False。side 引數功能相同，並且更推薦使用。如果將 side 設定為 “left” 同時此引數為 True，則會出錯。

• side (str, optional) – 功能與 right 相同，但更推薦使用。“left” 對應 right 的 False，“right” 對應 right 的 True。如果將此引數設定為 “left” 同時 right 為 True，則會出錯。預設值為 None。

• out (Tensor, optional) – 輸出張量，如果提供，其大小必須與 values 相同。

• sorter (LongTensor, optional) – 如果提供，一個與未排序的 sorted_sequence 形狀匹配的張量，包含在最內層維度上按升序對其進行排序的索引序列

示例

>>> 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]])