torch.bucketize

torch.bucketize(input, boundaries, *, out_int32=False, right=False, out=None) → Tensor

返回每个值所属的桶的索引，其中桶的边界由 boundaries 设置。返回一个与 input 大小相同的新的张量。如果 right 为 False（默认），则左边界为开区间。请注意，这与 numpy.digitize 的行为相反。更正式地说，返回的索引满足以下规则：

right	返回的索引满足条件
False	boundaries[i-1] < input[m][n]...[l][x] <= boundaries[i]
True	boundaries[i-1] <= input[m][n]...[l][x] < boundaries[i]

参数:

• 输入（Tensor 或 Scalar）- N 维张量或包含搜索值（s）的标量。

• 边界（Tensor）- 1 维张量，必须包含严格递增的序列，否则返回值未定义。

关键字参数:

• out_int32（bool，可选）- 指示输出数据类型。如果为 True，则为 torch.int32，否则为 torch.int64。默认值为 False，即默认输出数据类型为 torch.int64。

• right（布尔值，可选）- 决定 boundaries 中值的处理行为。请参见上表。

• out（张量，可选）- 输出张量，如果提供，则必须与 input 的大小相同。

示例：

>>> boundaries = torch.tensor([1, 3, 5, 7, 9])
>>> boundaries
tensor([1, 3, 5, 7, 9])
>>> v = torch.tensor([[3, 6, 9], [3, 6, 9]])
>>> v
tensor([[3, 6, 9],
        [3, 6, 9]])
>>> torch.bucketize(v, boundaries)
tensor([[1, 3, 4],
        [1, 3, 4]])
>>> torch.bucketize(v, boundaries, right=True)
tensor([[2, 3, 5],
        [2, 3, 5]])