torch.sort()函数用法

torch.sort()函数用法

1.语法

torch.sort()按值升序沿给定维度对输入张量的元素进行排序。

torch.sort(input, dim=-1, descending=False, stable=False, *, out=None)

2.参数

• input(Tensor)：输入张量
• dim(int,optional)：要排序的维度。如果未给出dim，则选择输入的最后一个维度。
• descending(bool,optional)：控制排序顺序（升序或降序）。如果descending为True，则元素按值降序排序。
• stable(bool,optional)：使排序例程稳定，从而保证保留等效元素的顺序。如果stable为True，则排序例程变得稳定，保留等效元素的顺序。
• out(tuple,optional)：（Tensor，LongTensor）的输出元组，可选地将其用作输出缓冲区

3.返回

返回（值，索引）的元祖 (values, indices) ，其中值是排序后的值，索引是原始输入张量中元素的索引。

4.例子

x = torch.randn(3, 4)
print("Original data:\n", x)
# 降序排序
sorted, indices = torch.sort(x, descending=True)
# 排序后的值
print("Sort data:\n", sorted)
# 原始输入张量中元素的索引
print("indices:\n", indices)

参考

• torch.sort

😃😃😃