pytoch张量的极值和排序

通过pytorch获取张量（沿着某个维度）的最大值和最小值，以及这些值所在的位置。并对张量各个维度的值进行排序

import torch

a = torch.randn(3, 4)  # 建立一个3*4的张量
print(a)

print(torch.argmax(a, 0))  # 函数调用。返回的是沿着第0个维度，极大值所在位置，即每一列的最大值
print(a.argmax(0)) # 内置方法调用。返回的是沿着第0个维度，极大值所在位置，即每一列的最大值

print(torch.argmax(a, 1))  # 函数调用。返回的是沿着第1个维度，极大值所在位置，即每一行的最大值
print(a.argmax(1))  # 内置方法调用。返回的是沿着第1个维度，极大值所在位置，即每一行的最大值

print(torch.max(a, -1))  # 函数调用，返回的是沿着最后一个维度，包含极大值和极大值所在位置的元组
print(a.max(1)) # 内置方法调用，返回的是沿着第1个维度，包含极大值和极大值所在位置的元组

print(torch.min(a, -1))  # 函数调用，返回的是沿着最后一个维度，包含极小值和极小值所在位置的元组
print(a.min(0)) # 内置方法调用，返回的是沿着第0个维度，包含极小值和极小值所在位置的元组

print(a.sort(-1))  # 沿着最后一个维度排序（默认从小到大，如果要从大到小排序，需要设置参数，看下一行例子），返回排序后的张量和张量元素在该维度的原始位置
print(a.sort(-1, descending=True))

print(a.sort(0))  # 沿着第0个维度即对每一列排序（默认从小到大，如果要从大到小排序，需要设置参数，看下一行例子），返回排序后的张量和张量元素在该维度的原始位置
print(a.sort(0, descending=True))

Done!!!