How to convert a torch tensor to mi.Bool in a fast way?

[ciy405x]

I have a large GPU tensor composed of booleans and I want to convert it to the mi.Bool type. I've found a method that works, but it is rather slow:

batch_size = 10
num_of_choices = 4
choice = torch.multinomial(torch.rand(batch_size, num_of_choices).cuda(), 1)  # select a choice
active = mi.Bool((choice.squeeze() == 0).tolist())

I believe there are alternative approaches to improve upon this. Any advice would be greatly appreciated. To clarify, I cannot change the method of creating the input choice tensor (as shown above) to utilize Mitsuba's built-in types.
Thank you.

[merlinND]

Hello @ciy405x,

Does the following work for you?

choice = mi.Int32(choice.squeeze().to(torch.int32))
active = dr.eq(choice, 0)

[ciy405x]

Oh yes, this method is definitely faster than the one that I proposed! :)
I have another question regarding the Mitsuba library: when converting tensors, mi.Float(float_tensor) works fine for floating-point tensors of shape (n, ), but mi.Bool(bool_tensor) does not seem to work for boolean tensors. Is there a specific design philosophy that accounts for the difference in the conversion methods between floating-point and boolean tensors?

[wjakob]

To exchange tensors with other libraries, PyTorch uses an exchange standard called dlpack. The fact that you are running into issues with booelan-valued tensors is really a PyTorch limitation and not something related to Mitsuba or Dr.Jit. Here is a tiny reproducer of the issue:

>>> (torch.tensor([1,2,3])>2).__dlpack__()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/wjakob/.local/lib/python3.8/site-packages/torch/_tensor.py", line 1349, in __dlpack__
    return torch.to_dlpack(self)
RuntimeError: Bool type is not supported by dlpack

The error message that PyTorch gives is actually untrue. Boolean arrays are perfectly well supported by DLPack. I think that they just forgot to implement this case.

I opened a ticket on the PyTorch issue tracker: pytorch/pytorch#114673