torch.searchsorted

torch.searchsorted

官方文档链接:torch.searchsorted --- PyTorch 2.3 documentation

该函数用于在已排序的序列中查找要插入的值的位置,以保持序列的顺序,

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

参数如下,

  • sorted_sequence:这是一个N-D或1-D的张量,其中包含按最内部维度单调递增的序列。如果提供了sorter参数,则序列不需要按顺序排列

  • values:这是一个N-D张量或标量,包含要搜索的值

  • out_int32:这是一个可选参数,用于指示输出数据类型。如果为True,则输出数据类型为torch.int32,否则为torch.int64

  • right:这是一个可选参数,如果为False,则返回找到的第一个合适位置。如果为 True,则返回最后一个索引。如果找不到合适的索引,则对于非数值值(例如nan、inf),返回0,或者返回sorted_sequence内最内部维度的大小(超过最内部维度的最后一个索引)。如果为False,则获取每个值在sorted_sequence相应内部维度上的下限索引,如果为True,则获取上限索引。默认值为False

  • side:这是一个可选参数,"left" 对应于right为 False,"right" 对应于right为 True。如果将其设置为 "left",而right为 True,则会报错。默认值为None。

  • out:这是一个可选参数,输出张量,如果提供,则必须与 values 的大小相同

  • sorter:这是一个可选参数,如果提供,则是一个与未排序的sorted_sequence形状相匹配的张量,其中包含一个按最内部维度升序排列的索引序列

使用示例如下,

复制代码
sorted_sequence = torch.tensor([[1, 3, 5, 7, 9], [2, 4, 6, 8, 10]])
"""
tensor([[ 1,  3,  5,  7,  9],
        [ 2,  4,  6,  8, 10]])
"""

values = torch.tensor([[3, 6, 9], [3, 6, 9]])
"""
tensor([[3, 6, 9],
        [3, 6, 9]])
"""

torch.searchsorted(sorted_sequence, values)
"""
tensor([[1, 3, 4],
        [1, 2, 4]])
对于第一行 [3, 6, 9]:
数字3在第一行的sorted_sequence中的位置是索引1
数字6在第一行的sorted_sequence中的位置是索引3(6大于5而小于7,因此将6插入到索引3的位置时,能够使序列保持升序排序)
数字9在第一行的sorted_sequence中的位置是索引4
对于第二行 [3, 6, 9]:
数字3在第二行的sorted_sequence中的位置是索引1(3大于2而小于4,因此当索引为1时,不会改变序列的升序排序)
数字6在第二行的sorted_sequence中的位置是索引2
数字9在第二行的sorted_sequence中的位置是索引4(9大于8而小于10,因此当索引为4时,不会改变序列的升序排序)
"""

## 当side='right'时, 函数会返回每个值在对应行的sorted_sequence中的右侧插入位置索引
torch.searchsorted(sorted_sequence, values, side='right')
"""
tensor([[2, 3, 5],
        [1, 3, 4]])

对于第一行 [3, 6, 9]:
数字3在第一行的sorted_sequence中的右侧插入位置是索引2(数字3的右侧插入位置索引是2)
数字6在第一行的sorted_sequence中的右侧插入位置是索引3
数字9在第一行的sorted_sequence中的右侧插入位置是索引5(数字9的右侧插入位置索引是5)
对于第二行 [3, 6, 9]:
数字3在第二行的sorted_sequence中的右侧插入位置是索引1
数字6在第二行的sorted_sequence中的右侧插入位置是索引3(数字6的右侧插入位置索引是3)
数字9在第二行的sorted_sequence中的右侧插入位置是索引4
"""

sorted_sequence_1d = torch.tensor([1, 3, 5, 7, 9])
"""
tensor([1, 3, 5, 7, 9])
"""

torch.searchsorted(sorted_sequence_1d, values)
"""
tensor([[1, 3, 4],
        [1, 3, 4]])
"""
相关推荐
科研服务器mike_leeso4 分钟前
41 年 7 次转型!戴尔从 PC 到 AI 工厂的技术跃迁与组织重构
大数据·人工智能·机器学习
lisw056 分钟前
如何改善基于深度学习的场重构
深度学习·重构·软件工程
大千AI助手17 分钟前
机器学习模型评估指标AUC详解:从理论到实践
人工智能·机器学习·模型评估·roc·precision·recall·auc
2501_9139817826 分钟前
2025年智能家居无线数传设备品牌方案精选
大数据·人工智能·智能家居
不老刘30 分钟前
GitHub Spec-Kit:AI 时代的规范驱动开发工具
人工智能·github·spec-kit
mit6.82433 分钟前
[tile-lang] 张量核心 | 传统MMA->WGMMA | 底层自动选择优化
人工智能·chatgpt
csuzhucong33 分钟前
人类知识体系分类
人工智能·分类·数据挖掘
DisonTangor41 分钟前
Lumina-DiMOO:用于多模态生成与理解的全扩散大语言模型
人工智能·语言模型·自然语言处理·ai作画·aigc
golang学习记1 小时前
阿里又出手了,发布全新终端CLI工具,还支持VSCode
人工智能
机器之心1 小时前
具身智能迎来ImageNet时刻:RoboChallenge开放首个大规模真机基准测试集
人工智能·openai