[Relay][Frontend] Why Transpose before GatherND

class GatherND(OnnxOpConverter):
    """Operator converter for GatherND."""

    @classmethod
    def _impl_common(cls, data, indices, batch_dims=0):
        indices_dims = len(infer_shape(indices))
        indices_shape = infer_shape(indices)
        indices = _op.transpose(indices, axes=[-1] + list(range(indices_dims - 1)))
        index_rank = indices_shape[-1]
        return _op.gather_nd(data, indices, batch_dims, index_rank)

Observation

In python/tvm/relay/onnx.py, I found _op.transpose() was inserted before real _op.gather_nd operator. Actually, other OnnxOpConverters, like ScatterND have same or more relative insertions before the real ops.

In test_cases however, test_op_level3.py, input data are fed into relay.gather_nd directly.

Question

Is this _op.transpose necessary? What would happen to dataflow if we remove it?

Thanks!

I remember this transpose is necessary since our op follows MXNet convention. gather, gather_nd in MXNet have different expectation on indices argument compared to other framework.

The nn.dense & nn.batch_matmul have a same history I think. So recently I’ve added extra attrs to them to support the input tensor in transposed or non-transposed format.