We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Optimizer fails on shape inference error over native_batch_norm
The optimizer fails for the attach model (so dort fails as well). It was obtained with the latest onnx, onnxscript and torch nightly.

[dump3bug.zip](https://github.com/microsoft/onnxscript/files/15106272/dump3bug.zip)

To replicate:

```python
import onnx
from onnxscript import optimizer
onx = onnx.load(model)
optimized = optimizer.optimize(onx)
```

It is coming from the following graph module.

```
graph():
    %primals_7 : [num_users=1] = placeholder[target=primals_7]
    %primals_1 : [num_users=1] = placeholder[target=primals_1]
    %primals_2 : [num_users=1] = placeholder[target=primals_2]
    %primals_3 : [num_users=1] = placeholder[target=primals_3]
    %primals_4 : [num_users=1] = placeholder[target=primals_4]
    %primals_5 : [num_users=1] = placeholder[target=primals_5]
    %add : [num_users=2] = call_function[target=torch.ops.aten.add.Tensor](args = (%primals_7, %primals_1), kwargs = {})
    %_native_batch_norm_legit_no_training : [num_users=1] = call_function[target=torch.ops.aten._native_batch_norm_legit_no_training.default](args = (%add, %primals_2, %primals_3, %primals_4, %primals_5, 0.1, 1e-05), kwargs = {})
    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_no_training, 0), kwargs = {})
    return (add, getitem)
```

Error:

```
Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "check_model.py", line 43, in <module>
    optimized = optimizer.optimize(onx)
  File "onnxscript/onnxscript/optimizer/__init__.py", line 61, in optimize
    model = onnx.shape_inference.infer_shapes(
  File "onnx/onnx/shape_inference.py", line 46, in infer_shapes
    inferred_model_str = C.infer_shapes(
onnx.onnx_cpp2py_export.shape_inference.InferenceError: [ShapeInferenceError] Inference error(s): (op_type:_aten_native_batch_norm_inference_onnx, node name: _aten_native_batch_norm_inference_onnx_2): [ShapeInferenceError] Inferred shape and existing shape differ in dimension 0: (2) vs (0)
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `onnxscript/optimizer/__init__.py`
Content:
```
1 import logging
2 from typing import Any
3 
4 import onnx
5 import onnx.shape_inference
6 
7 from onnxscript import rewriter
8 from onnxscript.optimizer.constant_folding import fold_constants
9 from onnxscript.optimizer.copy_propagation import (
10     do_copy_propagation,
11     do_sequence_simplification,
12 )
13 from onnxscript.optimizer.remove_unused import remove_unused_nodes
14 from onnxscript.optimizer.remove_unused_function import remove_unused_functions
15 from onnxscript.optimizer.simple_function_folding import (
16     inline_functions_with_unused_outputs,
17     inline_simple_functions,
18 )
19 from onnxscript.rewriter import (
20     broadcast_to_matmul,
21     cast_constant_of_shape,
22     gemm_to_matmul_add,
23     no_op,
24 )
25 
26 logger = logging.getLogger(__name__)
27 
28 
29 def optimize(
30     model: onnx.ModelProto,
31     num_iterations: int = 2,
32     *,
33     onnx_shape_inference: bool = True,
34     stop_if_no_change: bool = True,
35     external_data_folder: str = "",
36     **kwargs: Any,
37 ) -> onnx.ModelProto:
38     """Optimize the model. Perform optimizations and clean-ups such as constant folding, dead code elimination, etc.
39 
40     Args:
41         model (onnx.ModelProto): The model to optimize.
42         num_iterations (int, optional): Number of iterations to perform.
43         onnx_shape_inference (bool, optional): Whether to perform onnx shape inference on the model.
44             Set this to False to turn off onnx shape inference, and rely on model carried shapes and types.
45             This is useful for models produced by PyTorch 2.2+ dynamo onnx exporter, where the model carries
46             the symbolic shapes recorded from dynamo tracing.
47         stop_if_no_change (bool, optional): Whether to stop if no change is detected.
48         external_data_folder (str, optional): The folder to store external data.
49         **kwargs: Additional keyword arguments. For BC purposes.
50     """
51     if kwargs.pop("function_aware_folding", None) is not None:
52         logger.warning(
53             "'function_aware_folding' is deprecated. 'optimize' now supports both fully inlined models and models with functions. "
54             "To achieve the same behavior as 'function_aware_folding=True' before, set 'onnx_shape_inference=False'. "
55             "This would turn off incremental onnx shape inference and rely on model carried shapes and types. "
56             "See 'onnx_shape_inference' for more details."
57         )
58     for _ in range(num_iterations):
59         if onnx_shape_inference:
60             if model.ByteSize() < 1024 * 1024 * 1024 * 2:
61                 model = onnx.shape_inference.infer_shapes(
62                     model, check_type=True, strict_mode=True, data_prop=True
63                 )
64             else:
65                 logger.warning(
66                     "The model size is too large for full model shape inference. "
67                     "Skipping this step."
68                 )
69 
70         inline_simple_functions(model)
71         modified = fold_constants(
72             model, external_data_folder, onnx_shape_inference=onnx_shape_inference
73         )
74 
75         remove_unused_nodes(model)
76         inline_simple_functions(model)
77         remove_unused_functions(model)
78         inline_functions_with_unused_outputs(model)
79         # NOTE: This is general rewrite rules
80         model = rewriter.rewrite(
81             model,
82             pattern_rewrite_rules=[
83                 *no_op.rules.rules,  # TODO: merge this rule into constant folding?
84                 *broadcast_to_matmul.rules.rules,
85                 gemm_to_matmul_add.rule,
86                 *cast_constant_of_shape.rules.rules,
87             ],
88         )
89         if stop_if_no_change and not modified:
90             logger.debug("Stopping after %d iterations.", _)
91             break
92 
93     for node in model.graph.node:
94         logger.debug("Node %s::%s name %s.", node.domain, node.op_type, node.name)
95 
96     for function in model.functions:
97         for node in function.node:
98             logger.debug(
99                 "Function %s::%s node %s::%s name %s.",
100                 function.domain,
101                 function.name,
102                 node.domain,
103                 node.op_type,
104                 node.name,
105             )
106 
107     # do_sequence_simplification(model)
108     return model
109 
110 
111 __all__ = [
112     "fold_constants",
113     "remove_unused_nodes",
114     "optimize",
115     "do_copy_propagation",
116     "do_sequence_simplification",
117 ]
118
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py
--- a/onnxscript/optimizer/__init__.py
+++ b/onnxscript/optimizer/__init__.py
@@ -58,8 +58,12 @@
     for _ in range(num_iterations):
         if onnx_shape_inference:
             if model.ByteSize() < 1024 * 1024 * 1024 * 2:
+                # NOTE: strict mode is disabled because it crashes on the models
+                # that have different shapes inferred from the model carried shapes.
+                # The case can be found in:
+                # https://github.com/microsoft/onnxscript/issues/1443
                 model = onnx.shape_inference.infer_shapes(
-                    model, check_type=True, strict_mode=True, data_prop=True
+                    model, check_type=True, strict_mode=False, data_prop=True
                 )
             else:
                 logger.warning(

{"golden_diff": "diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py\n--- a/onnxscript/optimizer/__init__.py\n+++ b/onnxscript/optimizer/__init__.py\n@@ -58,8 +58,12 @@\n     for _ in range(num_iterations):\n         if onnx_shape_inference:\n             if model.ByteSize() < 1024 * 1024 * 1024 * 2:\n+                # NOTE: strict mode is disabled because it crashes on the models\n+                # that have different shapes inferred from the model carried shapes.\n+                # The case can be found in:\n+                # https://github.com/microsoft/onnxscript/issues/1443\n                 model = onnx.shape_inference.infer_shapes(\n-                    model, check_type=True, strict_mode=True, data_prop=True\n+                    model, check_type=True, strict_mode=False, data_prop=True\n                 )\n             else:\n                 logger.warning(\n", "issue": "Optimizer fails on shape inference error over native_batch_norm\nThe optimizer fails for the attach model (so dort fails as well). It was obtained with the latest onnx, onnxscript and torch nightly.\r\n\r\n[dump3bug.zip](https://github.com/microsoft/onnxscript/files/15106272/dump3bug.zip)\r\n\r\nTo replicate:\r\n\r\n```python\r\nimport onnx\r\nfrom onnxscript import optimizer\r\nonx = onnx.load(model)\r\noptimized = optimizer.optimize(onx)\r\n```\r\n\r\nIt is coming from the following graph module.\r\n\r\n```\r\ngraph():\r\n    %primals_7 : [num_users=1] = placeholder[target=primals_7]\r\n    %primals_1 : [num_users=1] = placeholder[target=primals_1]\r\n    %primals_2 : [num_users=1] = placeholder[target=primals_2]\r\n    %primals_3 : [num_users=1] = placeholder[target=primals_3]\r\n    %primals_4 : [num_users=1] = placeholder[target=primals_4]\r\n    %primals_5 : [num_users=1] = placeholder[target=primals_5]\r\n    %add : [num_users=2] = call_function[target=torch.ops.aten.add.Tensor](args = (%primals_7, %primals_1), kwargs = {})\r\n    %_native_batch_norm_legit_no_training : [num_users=1] = call_function[target=torch.ops.aten._native_batch_norm_legit_no_training.default](args = (%add, %primals_2, %primals_3, %primals_4, %primals_5, 0.1, 1e-05), kwargs = {})\r\n    %getitem : [num_users=1] = call_function[target=operator.getitem](args = (%_native_batch_norm_legit_no_training, 0), kwargs = {})\r\n    return (add, getitem)\r\n```\r\n\r\nError:\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"/usr/lib/python3.10/runpy.py\", line 196, in _run_module_as_main\r\n    return _run_code(code, main_globals, None,\r\n  File \"/usr/lib/python3.10/runpy.py\", line 86, in _run_code\r\n    exec(code, run_globals)\r\n  File \"check_model.py\", line 43, in <module>\r\n    optimized = optimizer.optimize(onx)\r\n  File \"onnxscript/onnxscript/optimizer/__init__.py\", line 61, in optimize\r\n    model = onnx.shape_inference.infer_shapes(\r\n  File \"onnx/onnx/shape_inference.py\", line 46, in infer_shapes\r\n    inferred_model_str = C.infer_shapes(\r\nonnx.onnx_cpp2py_export.shape_inference.InferenceError: [ShapeInferenceError] Inference error(s): (op_type:_aten_native_batch_norm_inference_onnx, node name: _aten_native_batch_norm_inference_onnx_2): [ShapeInferenceError] Inferred shape and existing shape differ in dimension 0: (2) vs (0)\r\n```\n", "before_files": [{"content": "import logging\nfrom typing import Any\n\nimport onnx\nimport onnx.shape_inference\n\nfrom onnxscript import rewriter\nfrom onnxscript.optimizer.constant_folding import fold_constants\nfrom onnxscript.optimizer.copy_propagation import (\n    do_copy_propagation,\n    do_sequence_simplification,\n)\nfrom onnxscript.optimizer.remove_unused import remove_unused_nodes\nfrom onnxscript.optimizer.remove_unused_function import remove_unused_functions\nfrom onnxscript.optimizer.simple_function_folding import (\n    inline_functions_with_unused_outputs,\n    inline_simple_functions,\n)\nfrom onnxscript.rewriter import (\n    broadcast_to_matmul,\n    cast_constant_of_shape,\n    gemm_to_matmul_add,\n    no_op,\n)\n\nlogger = logging.getLogger(__name__)\n\n\ndef optimize(\n    model: onnx.ModelProto,\n    num_iterations: int = 2,\n    *,\n    onnx_shape_inference: bool = True,\n    stop_if_no_change: bool = True,\n    external_data_folder: str = \"\",\n    **kwargs: Any,\n) -> onnx.ModelProto:\n    \"\"\"Optimize the model. Perform optimizations and clean-ups such as constant folding, dead code elimination, etc.\n\n    Args:\n        model (onnx.ModelProto): The model to optimize.\n        num_iterations (int, optional): Number of iterations to perform.\n        onnx_shape_inference (bool, optional): Whether to perform onnx shape inference on the model.\n            Set this to False to turn off onnx shape inference, and rely on model carried shapes and types.\n            This is useful for models produced by PyTorch 2.2+ dynamo onnx exporter, where the model carries\n            the symbolic shapes recorded from dynamo tracing.\n        stop_if_no_change (bool, optional): Whether to stop if no change is detected.\n        external_data_folder (str, optional): The folder to store external data.\n        **kwargs: Additional keyword arguments. For BC purposes.\n    \"\"\"\n    if kwargs.pop(\"function_aware_folding\", None) is not None:\n        logger.warning(\n            \"'function_aware_folding' is deprecated. 'optimize' now supports both fully inlined models and models with functions. \"\n            \"To achieve the same behavior as 'function_aware_folding=True' before, set 'onnx_shape_inference=False'. \"\n            \"This would turn off incremental onnx shape inference and rely on model carried shapes and types. \"\n            \"See 'onnx_shape_inference' for more details.\"\n        )\n    for _ in range(num_iterations):\n        if onnx_shape_inference:\n            if model.ByteSize() < 1024 * 1024 * 1024 * 2:\n                model = onnx.shape_inference.infer_shapes(\n                    model, check_type=True, strict_mode=True, data_prop=True\n                )\n            else:\n                logger.warning(\n                    \"The model size is too large for full model shape inference. \"\n                    \"Skipping this step.\"\n                )\n\n        inline_simple_functions(model)\n        modified = fold_constants(\n            model, external_data_folder, onnx_shape_inference=onnx_shape_inference\n        )\n\n        remove_unused_nodes(model)\n        inline_simple_functions(model)\n        remove_unused_functions(model)\n        inline_functions_with_unused_outputs(model)\n        # NOTE: This is general rewrite rules\n        model = rewriter.rewrite(\n            model,\n            pattern_rewrite_rules=[\n                *no_op.rules.rules,  # TODO: merge this rule into constant folding?\n                *broadcast_to_matmul.rules.rules,\n                gemm_to_matmul_add.rule,\n                *cast_constant_of_shape.rules.rules,\n            ],\n        )\n        if stop_if_no_change and not modified:\n            logger.debug(\"Stopping after %d iterations.\", _)\n            break\n\n    for node in model.graph.node:\n        logger.debug(\"Node %s::%s name %s.\", node.domain, node.op_type, node.name)\n\n    for function in model.functions:\n        for node in function.node:\n            logger.debug(\n                \"Function %s::%s node %s::%s name %s.\",\n                function.domain,\n                function.name,\n                node.domain,\n                node.op_type,\n                node.name,\n            )\n\n    # do_sequence_simplification(model)\n    return model\n\n\n__all__ = [\n    \"fold_constants\",\n    \"remove_unused_nodes\",\n    \"optimize\",\n    \"do_copy_propagation\",\n    \"do_sequence_simplification\",\n]\n", "path": "onnxscript/optimizer/__init__.py"}], "after_files": [{"content": "import logging\nfrom typing import Any\n\nimport onnx\nimport onnx.shape_inference\n\nfrom onnxscript import rewriter\nfrom onnxscript.optimizer.constant_folding import fold_constants\nfrom onnxscript.optimizer.copy_propagation import (\n    do_copy_propagation,\n    do_sequence_simplification,\n)\nfrom onnxscript.optimizer.remove_unused import remove_unused_nodes\nfrom onnxscript.optimizer.remove_unused_function import remove_unused_functions\nfrom onnxscript.optimizer.simple_function_folding import (\n    inline_functions_with_unused_outputs,\n    inline_simple_functions,\n)\nfrom onnxscript.rewriter import (\n    broadcast_to_matmul,\n    cast_constant_of_shape,\n    gemm_to_matmul_add,\n    no_op,\n)\n\nlogger = logging.getLogger(__name__)\n\n\ndef optimize(\n    model: onnx.ModelProto,\n    num_iterations: int = 2,\n    *,\n    onnx_shape_inference: bool = True,\n    stop_if_no_change: bool = True,\n    external_data_folder: str = \"\",\n    **kwargs: Any,\n) -> onnx.ModelProto:\n    \"\"\"Optimize the model. Perform optimizations and clean-ups such as constant folding, dead code elimination, etc.\n\n    Args:\n        model (onnx.ModelProto): The model to optimize.\n        num_iterations (int, optional): Number of iterations to perform.\n        onnx_shape_inference (bool, optional): Whether to perform onnx shape inference on the model.\n            Set this to False to turn off onnx shape inference, and rely on model carried shapes and types.\n            This is useful for models produced by PyTorch 2.2+ dynamo onnx exporter, where the model carries\n            the symbolic shapes recorded from dynamo tracing.\n        stop_if_no_change (bool, optional): Whether to stop if no change is detected.\n        external_data_folder (str, optional): The folder to store external data.\n        **kwargs: Additional keyword arguments. For BC purposes.\n    \"\"\"\n    if kwargs.pop(\"function_aware_folding\", None) is not None:\n        logger.warning(\n            \"'function_aware_folding' is deprecated. 'optimize' now supports both fully inlined models and models with functions. \"\n            \"To achieve the same behavior as 'function_aware_folding=True' before, set 'onnx_shape_inference=False'. \"\n            \"This would turn off incremental onnx shape inference and rely on model carried shapes and types. \"\n            \"See 'onnx_shape_inference' for more details.\"\n        )\n    for _ in range(num_iterations):\n        if onnx_shape_inference:\n            if model.ByteSize() < 1024 * 1024 * 1024 * 2:\n                # NOTE: strict mode is disabled because it crashes on the models\n                # that have different shapes inferred from the model carried shapes.\n                # The case can be found in:\n                # https://github.com/microsoft/onnxscript/issues/1443\n                model = onnx.shape_inference.infer_shapes(\n                    model, check_type=True, strict_mode=False, data_prop=True\n                )\n            else:\n                logger.warning(\n                    \"The model size is too large for full model shape inference. \"\n                    \"Skipping this step.\"\n                )\n\n        inline_simple_functions(model)\n        modified = fold_constants(\n            model, external_data_folder, onnx_shape_inference=onnx_shape_inference\n        )\n\n        remove_unused_nodes(model)\n        inline_simple_functions(model)\n        remove_unused_functions(model)\n        inline_functions_with_unused_outputs(model)\n        # NOTE: This is general rewrite rules\n        model = rewriter.rewrite(\n            model,\n            pattern_rewrite_rules=[\n                *no_op.rules.rules,  # TODO: merge this rule into constant folding?\n                *broadcast_to_matmul.rules.rules,\n                gemm_to_matmul_add.rule,\n                *cast_constant_of_shape.rules.rules,\n            ],\n        )\n        if stop_if_no_change and not modified:\n            logger.debug(\"Stopping after %d iterations.\", _)\n            break\n\n    for node in model.graph.node:\n        logger.debug(\"Node %s::%s name %s.\", node.domain, node.op_type, node.name)\n\n    for function in model.functions:\n        for node in function.node:\n            logger.debug(\n                \"Function %s::%s node %s::%s name %s.\",\n                function.domain,\n                function.name,\n                node.domain,\n                node.op_type,\n                node.name,\n            )\n\n    # do_sequence_simplification(model)\n    return model\n\n\n__all__ = [\n    \"fold_constants\",\n    \"remove_unused_nodes\",\n    \"optimize\",\n    \"do_copy_propagation\",\n    \"do_sequence_simplification\",\n]\n", "path": "onnxscript/optimizer/__init__.py"}]}