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	返回的索引滿足
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]

參數

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

• values (Tensor 或 Scalar) – 包含搜尋值(s)的 N-D 張量或純量。

關鍵字參數

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

• right (bool, optional) – 如果為 False，則返回找到的第一個合適的位置。如果為 True，則返回最後一個這樣的索引。 如果找不到合適的索引，則對於非數值（例如 nan, inf）返回 0，或返回 sorted_sequence 中*最內層*維度的大小（超過*最內層*維度的最後一個索引）。 換句話說，如果為 False，則在 sorted_sequence 相應的*最內層*維度上獲取 values 中每個值的下界索引。 如果為 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]])