Skip to content

Commit 0cced88

Browse files
justinchubyJustin Chu
andauthored
Support None as op input in GraphBuilder (#2868)
## Summary Support \`None\` as a valid input to \`op.OpType()\` calls in the onnxscript builder. This is needed for ONNX ops with optional inputs (e.g., Gemm's optional C input). ## Changes - Add \`None\` to the \`VALUE_LIKE\` union type - Update \`_input_to_ir_value\` to return \`ir.Value | None\` and pass \`None\` through - Update \`_partition_inputs_attributes\` and \`call_op\` type signatures to accept \`None\` - Add tests for \`None\` inputs with standard and custom domain ops Also removed pylint because it is not doing much good nowadays. --------- Signed-off-by: Justin Chu <justinchuby@noreply.github.com> Co-authored-by: Justin Chu <justinchuby@noreply.github.com>
1 parent 847801c commit 0cced88

File tree

6 files changed

+40
-77
lines changed

6 files changed

+40
-77
lines changed

.lintrunner.toml

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -98,42 +98,6 @@ init_command = [
9898
]
9999
is_formatter = true
100100

101-
[[linter]]
102-
code = 'PYLINT'
103-
include_patterns = [
104-
'**/*.py',
105-
]
106-
exclude_patterns = [
107-
'docs/**',
108-
'examples/**',
109-
'onnxscript/_internal/converter_test.py',
110-
'onnxscript/optimizer/**', # FIXME
111-
'onnxscript/rewriter/**', # FIXME
112-
'tests/functions/**',
113-
'tests/models/**',
114-
'tests/onnx_backend_test_code/**',
115-
]
116-
command = [
117-
'python',
118-
'-m',
119-
'lintrunner_adapters',
120-
'run',
121-
'pylint_linter',
122-
'--rcfile=pyproject_pylint.toml',
123-
'--show-disable',
124-
'--',
125-
'@{{PATHSFILE}}'
126-
]
127-
init_command = [
128-
'python',
129-
'-m',
130-
'lintrunner_adapters',
131-
'run',
132-
'pip_init',
133-
'--dry-run={{DRYRUN}}',
134-
'--requirement=requirements/lintrunner/requirements.txt',
135-
]
136-
137101
[[linter]]
138102
code = 'EDITORCONFIG-CHECKER'
139103
include_patterns = ['**']

onnxscript/_internal/builder.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
Sequence[float],
3232
Sequence[bool],
3333
Sequence[str],
34+
None,
3435
]
3536

3637
# Mapping from Python scalar types to their default ONNX DataType,
@@ -258,7 +259,7 @@ def initializer(
258259

259260
def _input_to_ir_value(
260261
self, value: VALUE_LIKE, like_type: ir.Value | None = None
261-
) -> ir.Value:
262+
) -> ir.Value | None:
262263
"""Convert a permissible input (for a call to an op) into an ir.Value.
263264
264265
Permissible values include ir.Value as well as python constants that can be converted
@@ -267,6 +268,8 @@ def _input_to_ir_value(
267268
"""
268269
if isinstance(value, ir.Value):
269270
return value
271+
if value is None:
272+
return value
270273
dtype = (
271274
like_type.type.dtype
272275
if like_type is not None and like_type.type is not None
@@ -356,7 +359,7 @@ def _get_schema(
356359
def _partition_inputs_attributes(
357360
self,
358361
schema: onnx.defs.OpSchema | None,
359-
inputs: Sequence[ir.Value | ir.TensorProtocol],
362+
inputs: Sequence[ir.Value | ir.TensorProtocol | None],
360363
kwargs: dict[str, Any],
361364
) -> tuple[Sequence[ir.Value | ir.TensorProtocol], dict[str, Any]]:
362365
if schema is None:
@@ -504,7 +507,7 @@ def subgraph(
504507
def call_op(
505508
self,
506509
op_type: str,
507-
inputs: Sequence[ir.Value | ir.TensorProtocol],
510+
inputs: Sequence[ir.Value | ir.TensorProtocol | None],
508511
kwargs: dict[str, Any],
509512
):
510513
"""Create an ONNX node and add it to the graph, returning its output value(s)."""

onnxscript/_internal/builder_test.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -848,6 +848,39 @@ def add_mul(X, Y):
848848

849849
self.assertIn("does not match", str(cm.exception))
850850

851+
def test_none_input_is_passed_through(self):
852+
"""Test that None inputs are preserved as None in the node's inputs."""
853+
op, x, y = _create_builder_with_inputs()
854+
855+
# Gemm's third input (C) is optional; passing None should work
856+
result = op.Gemm(x, y, None, alpha=1.0)
857+
858+
nodes = list(op.builder.graph)
859+
self.assertEqual(len(nodes), 1)
860+
node = nodes[0]
861+
self.assertEqual(node.op_type, "Gemm")
862+
# The third input should be None (optional, omitted)
863+
self.assertEqual(len(list(node.inputs)), 3)
864+
self.assertIs(node.inputs[0], x)
865+
self.assertIs(node.inputs[1], y)
866+
self.assertIsNone(node.inputs[2])
867+
self.assertIsNotNone(result)
868+
869+
def test_none_input_with_custom_domain(self):
870+
"""Test that None inputs work with custom domain ops."""
871+
op, x, y = _create_builder_with_inputs()
872+
873+
result = op.CustomOp(x, None, y, _domain="com.custom")
874+
875+
nodes = list(op.builder.graph)
876+
self.assertEqual(len(nodes), 1)
877+
node = nodes[0]
878+
self.assertEqual(node.op_type, "CustomOp")
879+
self.assertIs(node.inputs[0], x)
880+
self.assertIsNone(node.inputs[1])
881+
self.assertIs(node.inputs[2], y)
882+
self.assertIsNotNone(result)
883+
851884

852885
class BuildSubgraphTest(unittest.TestCase):
853886
"""Tests for GraphBuilder.subgraph()."""

pyproject.toml

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,7 @@ module = [
8686
ignore_errors = true
8787

8888
[tool.pylint.messages_control]
89-
# NOTE: This list is for vscode. Add new disables in pyproject_pylint.toml for lintrunner
90-
# Exclude patterns should be modified in .lintrunner.toml
89+
# NOTE: This list is for vscode. Pylint is removed from CI
9190
disable = [
9291
"consider-using-from-import",
9392
"format",

pyproject_pylint.toml

Lines changed: 0 additions & 34 deletions
This file was deleted.

requirements/lintrunner/requirements.txt

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,5 @@ ruff==0.15.1
55
# MYPY
66
mypy==1.10.1
77
types-PyYAML==6.0.12.20250915
8-
# PYLINT
9-
pylint==3.3.9
108
# EDITORCONFIG-CHECKER
119
editorconfig-checker==3.4.1

0 commit comments

Comments
 (0)