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?
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.