topk() cannot be used with complex or bool type of a 0D(only cuda) or more D tensor but topk() can be used with complex or bool type of a 0D(only cpu) tensor

[hyperkai]

🐛 Describe the bug

It seems like topk() cannot be used with complex or bool type of a 0D(only cuda) or more D tensor according to the errors as shown below:

complex type of a 0D(only cuda) tensor:

import torch

my_tensor = torch.tensor(0.+0.j, device='cuda:0')

torch.topk(input=my_tensor, k=1) # RuntimeError: "topk_out_cuda" not implemented for 'ComplexFloat'

bool type of a 0D(only cuda) tensor:

import torch

my_tensor = torch.tensor(True, device='cuda:0')

torch.topk(input=my_tensor, k=1) # RuntimeError: "topk_out_cuda" not implemented for 'Bool'

complex type of a 1D tensor:

import torch

my_tensor = torch.tensor([0.+0.j, 1.+0.j, 2.+0.j], device='cpu')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_cpu" not implemented for 'ComplexFloat'

my_tensor = torch.tensor([0.+0.j, 1.+0.j, 2.+0.j], device='cuda:0')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_out_cuda" not implemented for 'ComplexFloat'

bool type of a 1D tensor:

import torch

my_tensor = torch.tensor([True, False, True], device='cpu')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_cpu" not implemented for 'Bool'

my_tensor = torch.tensor([True, False, True], device='cuda:0')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_out_cuda" not implemented for 'Bool'

complex type of a 2D tensor:

import torch

my_tensor = torch.tensor([[0.+0.j, 1.+0.j], [2.+0.j, 3.+0.j]], device='cpu')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_cpu" not implemented for 'ComplexFloat'

my_tensor = torch.tensor([[0.+0.j, 1.+0.j], [2.+0.j, 3.+0.j]], device='cuda:0')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_out_cuda" not implemented for 'ComplexFloat'

bool type of a 2D tensor:

import torch

my_tensor = torch.tensor([[True, False], [False, True]],  device='cpu')

torch.topk(input=my_tensor, k=2)  # RuntimeError: "topk_cpu" not implemented for 'Bool'

my_tensor = torch.tensor([[True, False], [False, True]],  device='cuda:0')

torch.topk(input=my_tensor, k=2) # RuntimeError: "topk_out_cuda" not implemented for 'Bool'

But topk() can be used with complex or bool type of a 0D(only cpu) tensor as shown below:

complex type of a 0D(only cpu) tensor:

import torch

my_tensor = torch.tensor(0.+0.j, device='cpu')

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

my_tensor = torch.tensor(0.+0.j, device='cuda:0')

torch.topk(input=my_tensor, k=1) # RuntimeError: "topk_out_cuda" not implemented for 'ComplexFloat'

bool type of a 0D tensor:

import torch

my_tensor = torch.tensor(True, device='cpu')

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

Versions

import torch

torch.__version__ # 2.2.1+cu121

[malfet]

We would accept a PR that errors out even on empty, single element tensors if dtype are complex or boolean types