Skip to content

Commit cc1edfa

Browse files
gramalingamCopilot
andcommitted
Skip shape inference when outputs already have static type/shape
When constant propagation has already determined the type and static shape of all node outputs, skip the expensive ONNX shape inference call. Only skip when shapes are fully static (all dims are concrete integers); dynamic/symbolic dims still benefit from inference. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent df5b723 commit cc1edfa

1 file changed

Lines changed: 11 additions & 0 deletions

File tree

onnxscript/_internal/_inference.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,18 @@ def get_type(index: int, value: ir.Value) -> onnx.TypeProto:
9393
output.type = ir.serde.deserialize_type_proto_for_type(inferred_type)
9494

9595

96+
def _has_complete_type_and_shape(output: ir.Value) -> bool:
97+
"""Check if an output already has fully determined type and static shape."""
98+
if output.type is None or output.shape is None:
99+
return False
100+
return output.shape.is_static()
101+
102+
96103
def infer_outputs(node: ir.Node) -> None:
104+
# Skip inference if all outputs already have type and fully static shape
105+
# (e.g., from constant propagation setting const_value).
106+
if all(_has_complete_type_and_shape(output) for output in node.outputs):
107+
return
97108
try:
98109
_do_onnx_inference(node)
99110
except Exception as e: # pylint: disable=broad-exception-caught

0 commit comments

Comments
 (0)