Please implement where_cpu and where_cuda for BoolTensors

[sshish]

🚀 Feature

mycudabooltensor.where(...)
returns:
RuntimeError: "where_cuda" not implemented for 'Bool'

Motivation

In order to do masked operations on bool tensors, I have to cast those to byte tensors first, which looks unelegant.

cc @izdeby

[izdeby]

Done

[TheodoreZhao]

There is also the same kind of problem on CPU.
RuntimeError: "where_cpu" not implemented for 'Bool'

[izdeby]

Reopening in order to fix the CPU path

[izdeby]

@TheodoreZhao, can you provide repro code? This should work.
Tested via

import torch
a = torch.tensor([True, False], device='cuda')
res = torch.where(a > 0)
print(res)


a = torch.tensor([True, False], device='cpu')
res = torch.where(a > 0)
print(res)

Output:
(tensor([0], device='cuda:0'),)
(tensor([0]),)

[TheodoreZhao]

@izdeby It looks the types of x and y in torch.where(condition, x, y) cannot be BoolTensor.
https://pytorch.org/docs/stable/torch.html#torch.where

import torch
a = torch.tensor([True, False], device='cuda')
x = torch.tensor([True, True], device='cuda')
y = torch.tensor([False, False], device='cuda')
res = torch.where(a, x, y)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module> 
RuntimeError: "where_cuda" not implemented for 'Bool'

import torch
a = torch.tensor([True, False], device='cpu')
x = torch.tensor([True, True], device='cpu')
y = torch.tensor([False, False], device='cpu')
res = torch.where(a, x, y)

Traceback (most recent call last):
  File "<stdin>", line 1, in <module> 
RuntimeError: "where_cpu" not implemented for 'Bool'

[izdeby]

Reopening while investigating

[spacegoing]

Any update on this?

[izdeby]

I have to fix this. Will do in next few days. sorry for the delay.

[genbei]

我必须解决这个问题。将在未来几天内完成。抱歉耽搁了。

Have you solved it now? I have also encountered this problem

[H-Huang]

Fixed in #47454

[Hananel-Hazan]

This issue still exist on PyTorch == 1.7 using cpu but seems to work find on cuda

On cpu:

import torch
a = torch.tensor([True, False], device='cpu')
x = torch.tensor([True, True], device='cpu')
y = torch.tensor([False, False], device='cpu')
res = torch.where(a, x, y)

File "<ipython-input-15-0d63a62ea507>", line 4, in <module>
res = torch.where(a, x, y)
RuntimeError: "where_cpu" not implemented for 'Bool'

On cuda:

import torch
a = torch.tensor([True, False], device='cuda')
x = torch.tensor([True, True], device='cuda')
y = torch.tensor([False, False], device='cuda')
res = torch.where(a, x, y)

Working as expected