torch-mlir
torch-mlir copied to clipboard
[Torch Dialect] emit aten.nonzero op, recompose aten.index.Tensor with aten.nonzero_numpy to aten.masked_select
In detail, this PR includes:
- emit
aten.nonzeroandaten.nonzero_numpyop; - add missing shape infer function for
aten.masked_selectop - add recomposition rule for
aten.index.Tensorop. When it's indices come fromaten.nonzero_numpyop and the shapes of associated tensor are exactly same, it's equivalent to aaten.masked_selectop. But for now, no backend supportsaten.masked_selectoraten.nonzero_numpyop, this PR only adds a MLIR test case.
PS: This PR relies on canonicalizer of prim.ListConstruct op which is introduced in PR #2306 .