From a6e2e196c2edddec3aec161e0b756d99cda5aab1 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Oct 2025 12:14:44 -0700 Subject: [PATCH 1/4] Assign names for None dims When onnx shape inference is run on symbolic input dims, it will not handle the dim name propagation and instead create a None. As long as we rely on the current version onnx shape inference there is not better information we can get. However, since in the optimizer we also have some custom shape propagator implemented (e.g. for Identity) that will propagate sym dims, we should encode the equivalents for those dimensions as much as possible. This PR assigns a string to all None dims produced by onnx shape inference, so that the string names can get propagated when possible by the optimizer. Signed-off-by: Justin Chu --- onnxscript/optimizer/_constant_folding.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 8317d2be6..868dcbc9b 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -974,6 +974,7 @@ def __init__( self._sizes: dict[str, int] = {} self._modified: bool = False self._state = OptimizerState() + self._unknown_dim_count = 0 self._reset() def _reset(self) -> None: @@ -982,6 +983,7 @@ def _reset(self) -> None: self._sizes = {} self._modified = False self._state = OptimizerState() + self._unknown_dim_count = 0 def _do_inference(self, node: ir.Node) -> None: output_types = {} @@ -1029,7 +1031,12 @@ def get_type(value: ir.Value) -> onnx.TypeProto | None: inferred_shape = ir.serde.deserialize_type_proto_for_shape( inferred_type ) - output.shape = _merge_shapes(output.shape, inferred_shape) + merged_shape = _merge_shapes(output.shape, inferred_shape) + assert merged_shape is not None + output.shape = merged_shape + for i in range(len(merged_shape)): + if merged_shape.is_unknown_dim(i): + merged_shape[i] = ir.SymbolicDim(self._new_unknown_dim_name()) output.type = ir.serde.deserialize_type_proto_for_type(inferred_type) except Exception as e: logger.debug( @@ -1038,6 +1045,11 @@ def get_type(value: ir.Value) -> onnx.TypeProto | None: e, ) + def _new_unknown_dim_name(self) -> str: + name = f"unknown_{self._unknown_dim_count}" + self._unknown_dim_count += 1 + return name + def new_constant(self, node: ir.Node, value) -> ir.Node | None: irvalue = node.outputs[0] if not isinstance(value, np.ndarray): From 9e4dbb0644b412b593de40afe21f9e92962f565f Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Oct 2025 12:17:40 -0700 Subject: [PATCH 2/4] update Signed-off-by: Justin Chu --- onnxscript/optimizer/_constant_folding.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 868dcbc9b..598c78d06 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -1032,11 +1032,14 @@ def get_type(value: ir.Value) -> onnx.TypeProto | None: inferred_type ) merged_shape = _merge_shapes(output.shape, inferred_shape) + + # Replace unknown dims with uniquely named symbolic dims assert merged_shape is not None - output.shape = merged_shape for i in range(len(merged_shape)): if merged_shape.is_unknown_dim(i): merged_shape[i] = ir.SymbolicDim(self._new_unknown_dim_name()) + + output.shape = merged_shape output.type = ir.serde.deserialize_type_proto_for_type(inferred_type) except Exception as e: logger.debug( From ff57d09edd1a281a8184109066a61bc0ba84e3f5 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Oct 2025 12:18:02 -0700 Subject: [PATCH 3/4] docs Signed-off-by: Justin Chu --- onnxscript/optimizer/_constant_folding.py | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 598c78d06..6332ff201 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -1049,6 +1049,7 @@ def get_type(value: ir.Value) -> onnx.TypeProto | None: ) def _new_unknown_dim_name(self) -> str: + """Generate a new unique name for an unknown (None) symbolic dimension.""" name = f"unknown_{self._unknown_dim_count}" self._unknown_dim_count += 1 return name From aad7773ad47e0a7f0c8dc71874f1c2563afe9b57 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Thu, 9 Oct 2025 12:18:24 -0700 Subject: [PATCH 4/4] docs Signed-off-by: Justin Chu --- onnxscript/optimizer/_constant_folding.py | 1 + 1 file changed, 1 insertion(+) diff --git a/onnxscript/optimizer/_constant_folding.py b/onnxscript/optimizer/_constant_folding.py index 6332ff201..8742faf58 100644 --- a/onnxscript/optimizer/_constant_folding.py +++ b/onnxscript/optimizer/_constant_folding.py @@ -974,6 +974,7 @@ def __init__( self._sizes: dict[str, int] = {} self._modified: bool = False self._state = OptimizerState() + # Count of unknown (None) symbolic dimensions seen so far for generating unique names self._unknown_dim_count = 0 self._reset()