How to export a simple model using List.__contains__ to ONNX

jit
1

When using torch.jit.script, the message shows that contains method is not supported.

This is a reduced part of my model, the function should be tagged with torch.jit.script because there’s a for loop using list.contains

And I want to export it to an onnx file but failed with the following output.

Code

from typing import List, Dict
import torch

x = torch.tensor([[59, 26, 32, 31, 58, 37, 12,  8,  8, 32, 27, 27, 35,  9,  3, 44, 22, 36,
                   22, 61, 51, 35, 15, 13, 14, 32, 22, 21,  9]], dtype=torch.long)

nums = [3, 4, 5, 6, 7, 8, 9, 14, 15, 16, 17, 18, 22, 23, 24, 25, 26, 27,
        28, 29, 30, 31, 37, 38, 39, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57]


@torch.jit.script
def batch(x, l: List[int]):
    for i in range(len(x)):
        for j in range(len(x[i])):
            if x[i, j] in l:
                x[i, j] *= 2
    return x


class Module1(torch.nn.Module):
    def forward(self, x):
        return batch(x, nums)


m1 = Module1()
print(m1(x))

torch.onnx.export(m1,
                  (x),
                  "2.onnx",
                  verbose=True,
                  input_names=["x"],
                  dynamic_axes={
                      "x": {
                          1: "frames",
                      },
                  },
                  opset_version=11,
                  )

Output

Traceback (most recent call last):
  File "E:\My Files\Projects\Python\test\test.py", line 28, in <module>
    torch.onnx.export(m1,
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\__init__.py", line 350, in export
    return utils.export(
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 163, in export
    _export(
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 1074, in _export
    graph, params_dict, torch_out = _model_to_graph(
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 731, in _model_to_graph
    graph = _optimize_graph(
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 308, in _optimize_graph
    graph = _C._jit_pass_onnx(graph, operator_export_type)
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\__init__.py", line 416, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 1401, in _run_symbolic_function
    return symbolic_fn(ctx, g, *inputs, **attrs)
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\symbolic_opset9.py", line 5064, in Loop
    torch._C._jit_pass_onnx_block(
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\__init__.py", line 416, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 1401, in _run_symbolic_function
    return symbolic_fn(ctx, g, *inputs, **attrs)
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\symbolic_opset9.py", line 5064, in Loop
    torch._C._jit_pass_onnx_block(
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\__init__.py", line 416, in _run_symbolic_function
    return utils._run_symbolic_function(*args, **kwargs)
  File "C:\CodeEnv\miniconda3\envs\dfs\lib\site-packages\torch\onnx\utils.py", line 1421, in _run_symbolic_function
    raise symbolic_registry.UnsupportedOperatorError(
torch.onnx.symbolic_registry.UnsupportedOperatorError: Exporting the operator ::__contains_ to ONNX opset version 11 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.

I’m using torch 1.12.1+cu113.