Attribute name 'dim' does not match provided name 'axis'

[justinchuby]

PyTorch ONNX Conversion Error Report

✅ Obtain model graph with `torch.export.export`
❌ Translate the graph into ONNX
⚪ Run `onnx.checker` on the ONNX model
⚪ Execute the model with ONNX Runtime
⚪ Validate model output accuracy

Error message:

Traceback (most recent call last):

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 368, in _call_op
    node := _construct_node(
            ^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 318, in _construct_node
    for attr in ir_convenience.convert_attributes(named_attrs)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/ir/_convenience.py", line 233, in convert_attributes
    attributes.append(convert_attribute(name, attr))
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/ir/_convenience.py", line 123, in convert_attribute
    raise ValueError(

ValueError: Attribute name 'dim' does not match provided name 'axis'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 412, in eval
    outputs = self._call_op(op_signature, named_inputs, named_attrs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 373, in _call_op
    raise errors.GraphConstructionError(

torch_onnx.errors.GraphConstructionError: Error constructing node for operator '::TopK'. named_inputs={'X': SymbolicTensor('arg0_1', type=Tensor(FLOAT), shape=[5], producer=None, index=None), 'K': SymbolicTensor('anonymous:13053915888', type=None, shape=None, producer=anonymous_node:13281009072, index=0)}, converted_named_inputs={'X': SymbolicTensor('arg0_1', type=Tensor(FLOAT), shape=[5], producer=None, index=None), 'K': SymbolicTensor('anonymous:13053915888', type=None, shape=None, producer=anonymous_node:13281009072, index=0)}, named_attrs={'axis': Attr('dim', INT, -1), 'largest': Attr('largest', INT, True), 'sorted': Attr('sorted', INT, True)}, opset=, op_signature=OpSignature(domain='', name='TopK', overload='', params=[Parameter(name='X', type_constraint=TypeConstraintParam(name='T', allowed_types={Tensor(INT64), Tensor(INT32), Tensor(DOUBLE), Tensor(UINT8), Tensor(UINT64), Tensor(INT16), Tensor(UINT16), Tensor(INT8), Tensor(FLOAT16), Tensor(FLOAT), Tensor(UINT32)}, description='Constrain input and output types to numeric tensors.'), required=True, variadic=False, default=_EMPTY_DEFAULT), Parameter(name='K', type_constraint=TypeConstraintParam(name='K', allowed_types={Tensor(INT64)}, description=''), required=True, variadic=False, default=_EMPTY_DEFAULT), AttributeParameter(name='axis', type=INT, required=False, default=AttrInt64('axis', -1)), AttributeParameter(name='largest', type=INT, required=False, default=AttrInt64('largest', 1)), AttributeParameter(name='sorted', type=INT, required=False, default=AttrInt64('sorted', 1))], outputs=[Parameter(name='Values', type_constraint=TypeConstraintParam(name='T', allowed_types={Tensor(INT64), Tensor(INT32), Tensor(DOUBLE), Tensor(UINT8), Tensor(UINT64), Tensor(INT16), Tensor(UINT16), Tensor(INT8), Tensor(FLOAT16), Tensor(FLOAT), Tensor(UINT32)}, description='Constrain input and output types to numeric tensors.'), required=True, variadic=False, default=_EMPTY_DEFAULT), Parameter(name='Indices', type_constraint=TypeConstraintParam(name='I', allowed_types={Tensor(INT64)}, description='Constrain index tensor to int64'), required=True, variadic=False, default=_EMPTY_DEFAULT)]).

The above exception was the direct cause of the following exception:

Traceback (most recent call last):

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 483, in eval_function
    return function.function(**named_inputs, **named_attrs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/function_libs/torch_lib/ops/core.py", line 8192, in aten_topk
    values, indices = op.TopK(self, k, axis=dim, largest=largest, sorted=sorted)
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/onnx_opset/_impl/opset11.py", line 3800, in TopK
    return op(
           ^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/values.py", line 300, in __call__
    return evaluator.default().eval(schema, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 417, in eval
    raise errors.GraphConstructionError(

torch_onnx.errors.GraphConstructionError: Error calling operator 'TopK' with args (SymbolicTensor('arg0_1', type=Tensor(FLOAT), shape=[5], producer=None, index=None), SymbolicTensor('anonymous:13053915888', type=None, shape=None, producer=anonymous_node:13281009072, index=0)) and kwargs {'axis': Attr('dim', INT, -1), 'largest': Attr('largest', INT, True), 'sorted': Attr('sorted', INT, True)}.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_core.py", line 379, in _handle_call_function_node_with_lowering
    outputs = onnx_function(*onnx_args, **onnx_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/values.py", line 528, in __call__
    return evaluator.default().eval_function(self, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_building.py", line 497, in eval_function
    raise errors.GraphConstructionError(

torch_onnx.errors.GraphConstructionError: Error calling function 'aten_topk' with args (SymbolicTensor('arg0_1', type=Tensor(FLOAT), shape=[5], producer=None, index=None), 3) and kwargs {}. The function is defined at '/Users/justinc/Documents/GitHub/torch-onnx/venv/lib/python3.11/site-packages/onnxscript/function_libs/torch_lib/ops/core.py:8182'.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_core.py", line 480, in _add_nodes
    _handle_call_function_node_with_lowering(

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_core.py", line 381, in _handle_call_function_node_with_lowering
    raise errors.GraphConstructionError(

torch_onnx.errors.GraphConstructionError: Error when calling function 'OnnxFunction(<function aten_topk at 0x13fe2c220>)' with args '[SymbolicTensor('arg0_1', type=Tensor(FLOAT), shape=[5], producer=None, index=None), 3]' and kwargs '{}'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_patch.py", line 196, in _torch_onnx_export
    ir_model = torch_onnx.exported_program_to_ir(program)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_core.py", line 619, in exported_program_to_ir
    values = _add_nodes(exported_program, model, lower=lower, registry=registry)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

  File "/Users/justinc/Documents/GitHub/torch-onnx/src/torch_onnx/_core.py", line 491, in _add_nodes
    raise errors.OnnxConversionError(

torch_onnx.errors.OnnxConversionError: Error when translating node %topk : [num_users=2] = call_function[target=torch.ops.aten.topk.default](args = (%arg0_1, 3), kwargs = {}). See the stack trace for more information.

Exported program:

ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, arg0_1: "f32[5]"):
            # File: /Users/justinc/Documents/GitHub/torch-onnx/tests/torch_tests/torch_onnx_test.py:3656 in forward, code: return torch.topk(x, 3)
            topk = torch.ops.aten.topk.default(arg0_1, 3);  arg0_1 = None
            getitem: "f32[3]" = topk[0]
            getitem_1: "i64[3]" = topk[1];  topk = None
            return (getitem, getitem_1)

Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=<InputKind.USER_INPUT: 1>, arg=TensorArgument(name='arg0_1'), target=None, persistent=None)], output_specs=[OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='getitem'), target=None), OutputSpec(kind=<OutputKind.USER_OUTPUT: 1>, arg=TensorArgument(name='getitem_1'), target=None)])
Range constraints: {}

Analysis

PyTorch ONNX Conversion Analysis

Model Information

The model has 0 parameters and 0 buffers (non-trainable parameters). Number of parameters per dtype:

defaultdict(<class 'int'>, {})

Number of buffers per dtype:

defaultdict(<class 'int'>, {})

Inputs:

• arg0_1: TensorMetadata(shape=torch.Size([5]), dtype=torch.float32, requires_grad=True, stride=(1,), memory_format=torch.contiguous_format, is_quantized=False, qparams={})

Outputs:

• getitem: TensorMetadata(shape=torch.Size([3]), dtype=torch.float32, requires_grad=False, stride=(1,), memory_format=torch.contiguous_format, is_quantized=False, qparams={})
• getitem_1: TensorMetadata(shape=torch.Size([3]), dtype=torch.int64, requires_grad=False, stride=(1,), memory_format=torch.contiguous_format, is_quantized=False, qparams={})

The FX graph has 5 nodes in total. Number of FX nodes per op:

• placeholder: 1
• call_function: 3
• output: 1

Of the call_function nodes, the counts of operators used are:

• <built-in function getitem>: 2
• aten.topk.default: 1

ONNX Conversion Information

All operators in the model have registered ONNX decompositions.

[justinchuby]

Convert attr needs to copy the attribute if it is taken from the function default. Attr objects are used as values.