Repository: microsoft/onnxscript

Issue:
Wrong shape of mean and variance in _aten_native_batch_norm_inference_onnx
In the implementation of `_aten_native_batch_norm_inference_onnx()` the shapes of mean and variance seem wrong. They should be length-C vectors according to the ONNX spec but the current aten definitions result in scalars.
https://github.com/microsoft/onnxscript/blob/main/onnxscript/function_libs/torch_lib/ops/core.py#L5657-L5658
```
    empty_mean = op.CastLike(op.Shape(input, start=0, end=0), norm)
    empty_var = op.CastLike(op.Shape(input, start=0, end=0), norm)
```

Additionally, in some of the dynamo exported models, I see that the `running_mean` and `running_var` inputs also have wrong data type:
![image](https://github.com/microsoft/onnxscript/assets/99693700/d77b1007-cbf8-463b-b470-8ce8d734ffc1)
![image](https://github.com/microsoft/onnxscript/assets/99693700/fae68f3f-27ec-4698-813b-0097334ca0cc)



Discussion:
You are absolutely right :D. We did not implement it because the model did not use those outputs at the time. Contributions are welcome. Strange it is "true" though. I wonder what argument PyTorch supplied as running_mean
I think in most cases (where these outputs are unused), they are optimized away by function aware onnx rewriter.
It might indeed prevent the use of onnx shape inference before rewriting, since the phony outputs would still exist in the graph and lead to errors.
We need to look into that `running: boolean`
I looked deeper and realized the shapes are expected - they follow what PyTorch produces. When training is False, they are empty.
> I looked deeper and realized the shapes are expected - they follow what PyTorch produces. When training is False, they are empty.

Edit: but that's true only for CPU. We should ignore cpu in this case and follow the GPU behavior. @BowenBao was that what you decided on?

Resolving patch:
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
--- a/onnxscript/function_libs/torch_lib/ops/core.py
+++ b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -5757,8 +5757,9 @@ def _aten_native_batch_norm_inference_onnx(
     # We use CUDA's output here
     invstd = op.Div(1.0, op.Sqrt(running_var + eps))
     # https://github.com/pytorch/pytorch/blob/a44f8894fa6d973693aab44a3dda079a168b05c1/torch/_decomp/decompositions.py#L1475
-    # TODO(justinchuby): Make sure the output types are correct
-    return norm, running_mean, invstd, running_mean, running_var
+    running_mean_fp32 = op.Cast(running_mean, to=FLOAT.dtype)
+    invstd = op.Cast(invstd, to=FLOAT.dtype)
+    return norm, running_mean_fp32, invstd, running_mean, running_var
 
 
 # TODO: This op is using duplicated code from aten_native_batch_norm,


Test patch:
diff --git a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
--- a/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
+++ b/onnxscript/tests/function_libs/torch_lib/ops_test_data.py
@@ -1849,7 +1849,10 @@ def _where_input_wrangler(
         reason="native_batch_norm outputs different dtypes on CPU and CUDA. Our implematation is based on that for CUDA",
     ),
     TorchLibOpInfo(
-        "ops.aten._native_batch_norm_legit", core_ops.aten_native_batch_norm, trace_only=True
+        "ops.aten._native_batch_norm_legit",
+        core_ops.aten_native_batch_norm,
+        trace_only=True,
+        tolerance={torch.float16: (1e-2, 7e-3)},
     ).skip(
         device_type="cpu",
         matcher=lambda sample: sample.kwargs.get("training") is False,
@@ -1864,10 +1867,19 @@ def _where_input_wrangler(
         "ops.aten._native_batch_norm_legit_functional",
         core_ops.aten__native_batch_norm_legit_functional,
         trace_only=True,
-    ).skip(
+        tolerance={torch.float16: (1e-2, 7e-3)},
+    )
+    .skip(
         device_type="cpu",
         matcher=lambda sample: sample.kwargs.get("training") is False,
         reason="native_batch_norm outputs different results on CPU and CUDA when training is False. Our implematation is based on that for CUDA",
+    )
+    .skip(
+        dtypes=(torch.float16,),
+        device_type="cuda",
+        matcher=lambda sample: sample.kwargs.get("training") is True,
+        test_class_name="TestOutputConsistencyEager",
+        reason="fixme: output 4 (new_running_var) does not match the gpu output sometimes",
     ),
     TorchLibOpInfo(
         "ops.aten.native_group_norm",