[Python] Export a PyTorch custom hash table lookup OP to ONNX

I have implemented a PyTorch OP which accepts a torch.int64 tensor and outputs another torch.int64 tensor by looking up a hash table with predefined key-value pairs.

The torch part is implemented like this, where table is a wrapper around a C++ hash table:

class TableLookup(torch.autograd.Function):
@staticmethod
def forward(ctx, table, x):
y = table.lookup(x)
return y

@staticmethod
def symbolic(g, table, x):
raise NotImplementedError


The forward pass works nicely, backward isn't needed since x and y are integral tensors

My question is how to implement the symbolic so that torch could correctly export this OP to ONNX format.

I just checked ONNX specification, the version 2 of the "ai.onnx.ml.LabelEncoder" OP could probably be used as a table lookup OP.

Continue reading...