pytorch
✨[Feature] Expand index converter support to boolean masks
Current index converter only supports int type while cannot handle bool mask indices.
We need to expand index support to boolean masks, since this is a valid usage in Torch and we already support something very similar in aten.where.
Reported in the bug: #2480
I think for this converter, the output shape is dynamic because it depends on the mask. It's similar to a bug reported in #2516. So let's wait for getting TRT bug fixed.
I think for this converter, the output shape is dynamic because it depends on the mask. It's similar to a bug reported in #2516. So let's wait for getting TRT bug fixed.
@zewenli98 One thing to note in the case of x[mask] is that mask is a boolean tensor. And in the index converter implementation, https://github.com/pytorch/TensorRT/blob/main/py/torch_tensorrt/dynamo/conversion/impl/select.py#L130
the indices tensor received would be a boolean mask. This boolean mask should be converted to indices (where the non zero valued indices should be extracted. I'm not sure if TRT has such an operator)
