Repository: microsoft/onnxscript

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


Discussion:
@titaiwangms Do you have an idea? Looks like it is related to https://github.com/microsoft/onnxscript/blame/0d98619dee85025f8fb110864607f6f477c3d8ae/onnxscript/function_libs/torch_lib/ops/core.py#L5625
I can take a look after I done my work on hands,
@xadupre I think optimizer should be already applied in torch-nightly? https://github.com/pytorch/pytorch/blob/d5182bb75bbc109cb327212e7205981fbf72cb5e/torch/onnx/_internal/exporter.py#L1274

Are you writing new optimization? Just trying to understand the usage here.

Specifically, if we try:

```python
import onnx
from onnxscript import optimizer
from onnxscript.rewriter import onnxruntime as ort_rewriter

onx = onnx.load("dump3bug.onnx")
onnx.checker.check_model(onx, full_check=True)
optimized = optimizer.optimize(onx)
```

The same error is spotted by checker.
It does. I tweaked torch code in onnxruntime.py to get the model before it gets optimized to know of the error happens before optimization or after. It is after.
Could you update the model to the one before optimizer?
I'll check again but it should be the one before optimizer.
I mean the model in zip doesn't pass onnx.checker.check_model(model, full_check=True). That's why it gets the error message from https://github.com/microsoft/onnxscript/blob/2a78f61b5d4e31bd591cc020211aa903bc425270/onnxscript/optimizer/__init__.py#L61

It's not even hitting the constant folding and general rewriter yet it seems.
I wonder if we should put onnx.checker to guard the models generated from converter/dort. Or we already did?
> I mean the model in zip doesn't pass onnx.checker.check_model(model, full_check=True).

True ... tried it, and this seems to fail
I would not call onnx.checker. The converter may introduce nodes coming from domain *com.microsoft*. I created PR #1467 to replicate the issue.
So I think there are two issues here. The first one is that if we don't want to make sure our models passing checker before feeding to optimizer, we should turn off strict_mode in ONNX shape type inference inside optimizer, since they are basically the same. I will submit a PR for this to unblock this model.

The other issue is that, in torchlib, we respect PyTorch native_batch_norm CUDA to accept size=0 outputs in index=1 and 2 ([here](https://github.com/microsoft/onnxscript/blob/8dba367fb000e3696c79b618638861b5cdf759dc/onnxscript/function_libs/torch_lib/ops/core.py#L5797-L5804)), which is originated from [PyTorch code](https://github.com/pytorch/pytorch/blob/a44f8894fa6d973693aab44a3dda079a168b05c1/torch/_decomp/decompositions.py#L1452-L1457). That's why in the error message saying the existing shape is 0. However, ONNX shape type inference infers this as 2. @justinchuby @xiaowuhu @gramalingam any suggestion on this?

Resolving patch:
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 @@ def optimize(
     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(


Test patch: