From 65950e26e7559ea846a1b3d6ac7a670fd0b98111 Mon Sep 17 00:00:00 2001 From: Theodore Cooper Date: Sun, 12 Jul 2026 19:10:41 +0800 Subject: [PATCH] Fix optional input positions in converter --- onnxscript/_internal/converter.py | 7 ++- onnxscript/_internal/converter_test.py | 12 ++++ onnxscript/_internal/param_manipulation.py | 20 ++++++- .../_internal/param_manipulation_test.py | 58 +++++++++++++++++++ 4 files changed, 92 insertions(+), 5 deletions(-) diff --git a/onnxscript/_internal/converter.py b/onnxscript/_internal/converter.py index c2f6b0fb63..86536eec19 100644 --- a/onnxscript/_internal/converter.py +++ b/onnxscript/_internal/converter.py @@ -619,10 +619,13 @@ def _translate_expr(self, node: ast.AST, target: PreferredName | None = None) -> result = self._generate_unique_name(target) return self._emit1([result], callee, args, attrs) - def _translate_opt_expr(self, node: ast.expr) -> ir.Value | None: + def _translate_opt_expr(self, node: ast.expr | None) -> ir.Value | None: """Translation of an expression where "None" is permitted (eg., for an optional argument). - None is represented as a Constant in Python 3.9+. + Parsed None is represented as a Constant in Python 3.9+, while argument + normalization may insert None directly for an omitted input. """ + if node is None: + return None if isinstance(node, ast.Constant) and (node.value is None): return None return self._translate_expr(node) diff --git a/onnxscript/_internal/converter_test.py b/onnxscript/_internal/converter_test.py index c1fe276ef5..d52784ba54 100644 --- a/onnxscript/_internal/converter_test.py +++ b/onnxscript/_internal/converter_test.py @@ -575,6 +575,18 @@ def keyword(X, shape): onnxscript.testing.assert_isomorphic_function(positional, keyword) + def test_keyword_input_preserves_optional_input_positions(self): + @script() + def resize(X: FLOAT[3, 4]) -> FLOAT[2, 2]: + return op.Resize(X, sizes=[2, 2]) + + model = resize.to_model_proto() + resize_node = next(node for node in model.graph.node if node.op_type == "Resize") + + self.assertEqual(list(resize_node.input[1:3]), ["", ""]) + self.assertEqual(len(resize_node.input), 4) + create_cpu_inference_session(model.SerializeToString()) + def test_none_as_input_for_op_with_no_schema(self): """Test conversion of None as an input value in a call to an op with no known schema.""" diff --git a/onnxscript/_internal/param_manipulation.py b/onnxscript/_internal/param_manipulation.py index c75d42504b..443a7c4e81 100644 --- a/onnxscript/_internal/param_manipulation.py +++ b/onnxscript/_internal/param_manipulation.py @@ -52,6 +52,14 @@ def separate_input_attributes_from_arguments( onnx_inputs = [] onnx_attributes = collections.OrderedDict() has_variadic = False + omitted_optional_inputs = 0 + + def append_input(value: Any) -> None: + nonlocal omitted_optional_inputs + # Omitted inputs need placeholders only when a later input is present. + onnx_inputs.extend([None] * omitted_optional_inputs) + omitted_optional_inputs = 0 + onnx_inputs.append(value) for i, param in enumerate(op_signature.params): is_input = param.is_param() @@ -60,17 +68,21 @@ def separate_input_attributes_from_arguments( if is_variadic: has_variadic = True # Exhaust all remaining args - onnx_inputs.extend(args[i:]) + variadic_args = args[i:] + if variadic_args: + onnx_inputs.extend([None] * omitted_optional_inputs) + omitted_optional_inputs = 0 + onnx_inputs.extend(variadic_args) args = [] continue if i < len(args): if is_input: - onnx_inputs.append(args[i]) + append_input(args[i]) else: onnx_attributes[param.name] = args[i] elif param.name in kwargs: if is_input: - onnx_inputs.append(kwargs[param.name]) + append_input(kwargs[param.name]) else: onnx_attributes[param.name] = kwargs[param.name] elif isinstance(param, ir.schemas.AttributeParameter) and param.has_default(): @@ -80,6 +92,8 @@ def separate_input_attributes_from_arguments( onnx_attributes[param.name] = param.default.value elif param.required: raise TypeError(f"Required input '{param}' was not provided") + elif is_input: + omitted_optional_inputs += 1 if not allow_extra_args and not has_variadic and len(args) > len(op_signature.params): raise TypeError( diff --git a/onnxscript/_internal/param_manipulation_test.py b/onnxscript/_internal/param_manipulation_test.py index 892ecd28e7..30e5b6c7d1 100644 --- a/onnxscript/_internal/param_manipulation_test.py +++ b/onnxscript/_internal/param_manipulation_test.py @@ -223,6 +223,64 @@ def test_it_does_not_fill_default_when_fill_defaults_is_false( self.assertEqual(inputs, [TEST_INPUT]) self.assertEqual(attributes, collections.OrderedDict([("b", 42)])) + @parameterized.parameterized.expand( + [ + ( + "later_optional_input", + (TEST_INPUT,), + {"d": "LAST_INPUT"}, + [TEST_INPUT, None, None, "LAST_INPUT"], + ), + ( + "middle_optional_input", + (TEST_INPUT,), + {"c": "MIDDLE_INPUT"}, + [TEST_INPUT, None, "MIDDLE_INPUT"], + ), + ( + "trailing_optional_inputs_omitted", + (TEST_INPUT,), + {}, + [TEST_INPUT], + ), + ( + "explicit_none_is_preserved", + (TEST_INPUT, None), + {}, + [TEST_INPUT, None], + ), + ] + ) + def test_optional_input_positions_are_preserved(self, _, args, kwargs, expected_inputs): + type_constraint = ir.schemas.TypeConstraintParam.any_tensor("T") + op_signature = ir.schemas.OpSignature( + domain="", + name="TestOp", + overload="", + params=[ + ir.schemas.Parameter( + name="a", type_constraint=type_constraint, required=True, variadic=False + ), + ir.schemas.Parameter( + name="b", type_constraint=type_constraint, required=False, variadic=False + ), + ir.schemas.Parameter( + name="c", type_constraint=type_constraint, required=False, variadic=False + ), + ir.schemas.Parameter( + name="d", type_constraint=type_constraint, required=False, variadic=False + ), + ], + outputs=[], + ) + + inputs, attributes = param_manipulation.separate_input_attributes_from_arguments( + op_signature, args, kwargs + ) + + self.assertEqual(inputs, expected_inputs) + self.assertEqual(attributes, collections.OrderedDict()) + @parameterized.parameterized.expand( [ (True, True),