The advanced indexing implementation of relay/topi is not able to handle cases that use None as one of the indexing argument.
My repro under test_forward_index
class Index2(Module):
def forward(self, x):
return x[None, [2, 2]]
input_data = torch.rand([3, 3, 3, 3]).float().cpu()
verify_model_with_input(Index2().eval(), [input_data])
Expected behavior
A tensor shaped [1, 2, 3, 3, 3] shall be returned.
Actual behavior
> raise Exception("warning unhandled case: {0}".format(type(expr)))
E Exception: warning unhandled case: <class 'NoneType'>
python/tvm/relay/expr_functor.py:76: Exception
Environment
Latest TVM main
cc: @masahi