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 <[email protected]>

Author: justinchuby

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):