DEV Community

Super Kai (Kazuya Ito)
Super Kai (Kazuya Ito)

Posted on • Updated on

kthvalue and topk in PyTorch

Buy Me a Coffee

*Memos:

kthvalue() can get two of the 0D or more D tensors of zero or more kth smallest elements and their indices from the 0D or more D tensor of zero or more elements as shown below:

*Memos:

  • kthvalue() can be used with torch or a tensor.
  • The 1st argument(input) with torch or using a tensor(Required-Type:tensor of int or float). *The 0D tensor of complex or bool works.
  • The 2nd argument with torch or the 1st argument with a tensor is k(Required-Type:int). *It must be greater than 0.
  • The 3rd argument with torch or the 2nd argument with a tensor is dim(Optional-Type:int).
  • The 4th argument with torch or the 3rd argument with a tensor is keepdim(Optional-Default:False-Type:bool): *Memos:
    • It must be used with dim.
    • My post explains keepdim argument.
  • There is out argument with torch(Optional-Default:None-Type:tuple(tensor, tensor) or list(tensor, tensor)): *Memos:
    • out= must be used.
    • My post explains out argument.
  • If there are the multiple same k th elements, one is returned nondeterministically.
  • An empty 2D or more D input tensor or tensor doesn't work if not setting dim with k=1.
  • An empty 1D input tesnor or tensor doesn't work even if setting dim with k=1.
import torch

my_tensor = torch.tensor([5, 1, 9, 7, 6, 8, 0, 5])

torch.kthvalue(input=my_tensor, k=3)
torch.kthvalue(input=my_tensor, k=3, dim=0)
torch.kthvalue(input=my_tensor, k=3, dim=-1)
# torch.return_types.kthvalue(
# values=tensor(5),
# indices=tensor(7))

torch.kthvalue(input=my_tensor, k=4)
torch.kthvalue(input=my_tensor, k=4, dim=0)
torch.kthvalue(input=my_tensor, k=4, dim=-1)
# torch.return_types.kthvalue(
# values=tensor(5),
# indices=tensor(0))

my_tensor = torch.tensor([[5, 1, 9, 7],
                          [6, 8, 0, 5]])
torch.kthvalue(input=my_tensor, k=3)
torch.kthvalue(input=my_tensor, k=3, dim=1)
torch.kthvalue(input=my_tensor, k=3, dim=-1)
# torch.return_types.kthvalue(
# values=tensor([7, 6]),
# indices=tensor([3, 0]))

my_tensor = torch.tensor([[5., 1., 9., 7.],
                          [6., 8., 0., 5.]])
torch.kthvalue(input=my_tensor, k=3)
# torch.return_types.kthvalue(
# values=tensor([7., 6.]),
# indices=tensor([3, 0]))

my_tensor = torch.tensor(5.+0.j)

torch.kthvalue(input=my_tensor, k=1)
# torch.return_types.kthvalue(
# values=tensor(5.+0.j),
# indices=tensor(0))

my_tensor = torch.tensor(True)

torch.kthvalue(input=my_tensor, k=1)
# torch.return_types.kthvalue(
# values=tensor(True),
# indices=tensor(0))

my_tensor = torch.tensor([])
my_tensor = torch.tensor([[]])
my_tensor = torch.tensor([[[]]])

torch.kthvalue(input=my_tensor, k=1) # Error

my_tensor = torch.tensor([])

torch.kthvalue(input=my_tensor, k=1, dim=0) # Error

my_tensor = torch.tensor([[]])

torch.kthvalue(input=my_tensor, k=1, dim=0)
# torch.return_types.kthvalue(
# values=tensor([]),
# indices=tensor([], dtype=torch.int64))

my_tensor = torch.tensor([[[]]])

torch.kthvalue(input=my_tensor, k=1, dim=0)
# torch.return_types.kthvalue(
# values=tensor([], size=(1, 0)),
# indices=tensor([], size=(1, 0), dtype=torch.int64))
Enter fullscreen mode Exit fullscreen mode

topk() can get two of the 0D or more D tensors of zero or more k largest or smallest elements and their indices from the 0D or more D tensor of zero or more elements as shown below:

*Memos:

  • topk() can be used with torch or a tensor.
  • The 1st argument(input) with torch or using a tensor(Required-Type:tensor of int or float). *complex or bool of a 0D tensor works on cpu.
  • The 2nd argument with torch or the 1st argument with a tensor is k(Required-Type:int). *It must be greater than or equal to 0.
  • The 3rd argument with torch or the 2nd argument with a tensor is dim(Optional-Type:int).
  • The 4th argument with torch or the 3rd argument with a tensor is largest(Optional-Default:True-Type:bool). *True gets zero or more largest elements while False gets zero or more smallest elements.
  • The 5th argument with torch or the 4th argument with a tensor is sorted(Optional-Default:True-Type:bool). *Sometimes, a return tensor is sorted with False but sometimes not so make it True if you want to definitely get a sorted tensor.
  • There is out argument with torch(Optional-Default:None-Type:tuple(tensor, tensor) or list(tensor, tensor)): *Memos:
    • out= must be used.
    • My post explains out argument.
  • If there are the multiple same k elements, one or more ones are returned nondeterministically.
  • An empty 2D or more D input tensor or tensor doesn't work if not setting dim with k=1.
  • An empty 1D input tensor or tensor doesn't work even if setting dim with k=1.
import torch

my_tensor = torch.tensor([5, 1, 9, 7, 6, 8, 0, 5])

torch.topk(input=my_tensor, k=3)
torch.topk(input=my_tensor, k=3, dim=0)
torch.topk(input=my_tensor, k=3, dim=-1)
# torch.return_types.topk(
# values=tensor([9, 8, 7]),
# indices=tensor([2, 5, 3]))

torch.topk(input=my_tensor, k=3, dim=0, largest=False)
# torch.return_types.topk(
# values=tensor([0, 1, 5]),
# indices=tensor([6, 1, 0]))

torch.topk(input=my_tensor, k=3, dim=0, largest=False, sorted=False)
# torch.return_types.topk(
# values=tensor([1, 0, 5]),
# indices=tensor([1, 6, 0]))

torch.topk(input=my_tensor, k=4)
torch.topk(input=my_tensor, k=4, dim=0)
torch.topk(input=my_tensor, k=4, dim=-1)
# torch.return_types.topk(
# values=tensor([9, 8, 7, 6]),
# indices=tensor([2, 5, 3, 4]))

torch.topk(input=my_tensor, k=4, dim=0, largest=False)
# torch.return_types.topk(
# values=tensor([0, 1, 5, 5]),
# indices=tensor([6, 1, 0, 7]))

torch.topk(input=my_tensor, k=4, dim=0, largest=False, sorted=False)
# torch.return_types.topk(
# values=tensor([1, 0, 5, 5]),
# indices=tensor([1, 6, 0, 7]))

my_tensor = torch.tensor([[5, 1, 9, 7],
                          [6, 8, 0, 5]])
torch.topk(input=my_tensor, k=3)
torch.topk(input=my_tensor, k=3, dim=1)
torch.topk(input=my_tensor, k=3, dim=-1)
# torch.return_types.topk(
# values=tensor([[9, 7, 5], [8, 6, 5]]),
# indices=tensor([[2, 3, 0], [1, 0, 3]]))

torch.topk(input=my_tensor, k=3, dim=1, largest=False)
# torch.return_types.topk(
# values=tensor([[1, 5, 7], [0, 5, 6]]),
# indices=tensor([[1, 0, 3], [2, 3, 0]]))

torch.topk(input=my_tensor, k=3, dim=1, largest=False, sorted=False)
# torch.return_types.topk(
# values=tensor([[1, 5, 7], [5, 0, 6]]),
# indices=tensor([[1, 0, 3], [3, 2, 0]]))

my_tensor = torch.tensor([[5., 1., 9., 7.],
                          [6., 8., 0., 5.]])
torch.topk(input=my_tensor, k=3)
# torch.return_types.topk(
# values=tensor([[9., 7., 5.],
#                [8., 6., 5.]]),
# indices=tensor([[2, 3, 0],
#                 [1, 0, 3]]))

my_tensor = torch.tensor(5.+0.j)

torch.topk(input=my_tensor, k=1)
# torch.return_types.topk(
# values=tensor(5.+0.j),
# indices=tensor(0))

my_tensor = torch.tensor(True)

torch.topk(input=my_tensor, k=1)
# torch.return_types.topk(
# values=tensor(True),
# indices=tensor(0))

my_tensor = torch.tensor([])
my_tensor = torch.tensor([[]])
my_tensor = torch.tensor([[[]]])

torch.topk(input=my_tensor, k=1) # Error

my_tensor = torch.tensor([])

torch.topk(input=my_tensor, k=1, dim=0) # Error

my_tensor = torch.tensor([[]])

torch.topk(input=my_tensor, k=1, dim=0)
# torch.return_types.topk(
# values=tensor([], size=(1, 0)),
# indices=tensor([], size=(1, 0), dtype=torch.int64))

my_tensor = torch.tensor([[[]]])

torch.topk(input=my_tensor, k=1, dim=0)
# torch.return_types.topk(
# values=tensor([], size=(1, 1, 0)),
# indices=tensor([], size=(1, 1, 0), dtype=torch.int64))
Enter fullscreen mode Exit fullscreen mode

Top comments (0)