How to do gathering on a tensor with two-dim indexing

1

Question

Hi,
Want to add symbolic func to a custom PyTorch op and export it to ONNX using existing ONNX ops. There is two-dim indexing operation. Have tried index_select, but not work. So could anyone take a look into this and help me with this?

Further information

Sample code

def my_custom_op(data, x_indices, y_indices):
    ## suppose this op is written in c++
    return data[x_indice, y_indices]

class MyCustomOp(torch.autograd.Function):
    
    @staticmethod
    def forward(ctx, data, x_indices, y_indices):
        return my_custom_op(data, x_indices, y_indices)

    @staticmethod
    def symbolic(g, data, x_indices, y_indices):
        from torch.onnx.symbolic_opset9 import index_select, transpose
        data_xs = index_select(g, data, 0, x_indices)
        ## don't know how to do this because index_select not work for this 
        # data_xs = transpose(g, data_xs, 0, 1)
        # data_ys = index_select(g, data_xs, 0, y_indices)
        return out

Thanks in advance.