onnxscript 0.1.0.dev20240123__tar.gz → 0.1.0.dev20240125__tar.gz

{onnxscript-0.1.0.dev20240123/onnxscript.egg-info → onnxscript-0.1.0.dev20240125}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20240123
+Version: 0.1.0.dev20240125
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>

{onnxscript-0.1.0.dev20240123 → onnxscript-0.1.0.dev20240125}/onnxscript/function_libs/torch_lib/ops/core.py

@@ -5566,10 +5566,16 @@ def aten_native_batch_norm(
     """native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor)"""
 
     if weight is None:  # Set to 1.0 as default
-        weight = op.Expand(op.Constant(value_floats=[1.0]), op.Shape(input, start=1, end=2))
+        weight = op.Expand(
+            op.CastLike(op.Constant(value_floats=[1.0]), input),
+            op.Shape(input, start=1, end=2),
+        )
 
     if bias is None:  # Set to 0.0 as default
-        bias = op.Expand(op.Constant(value_floats=[0.0]), op.Shape(input, start=1, end=2))
+        bias = op.Expand(
+            op.CastLike(op.Constant(value_floats=[0.0]), input),
+            op.Shape(input, start=1, end=2),
+        )
 
     axes = list(range(len(input.shape)))
     axes.pop(1)
@@ -5620,13 +5626,16 @@ def _aten_native_batch_norm_training_onnx(
         training_mode=training,
     )
     # Compute var and rstd
-    mean = op.ReduceMean(input, axes)
-    input_sub_mean = op.Sub(input, mean)
+    # Mean, var, and rstd computation and results are expected to be
+    # in higher precision when inputs are float16.
+    upcast_input = op.Cast(input, to=FLOAT.dtype)
+    mean = op.ReduceMean(upcast_input, axes)
+    input_sub_mean = op.Sub(upcast_input, mean)
     sqr = op.Mul(input_sub_mean, input_sub_mean)
     var = op.ReduceMean(sqr, axes, keepdims=False)
     rstd = op.Div(1.0, op.Sqrt(var + eps))
     # Get mean again with size = [1, C]
-    mean = op.ReduceMean(input, axes, keepdims=False)
+    mean = op.ReduceMean(upcast_input, axes, keepdims=False)
     return norm, mean, rstd
 
 
@@ -5735,13 +5744,16 @@ def _aten__native_batch_norm_training_functional_onnx(
         training_mode=training,
     )
     # Compute var and rstd
-    mean = op.ReduceMean(input, axes)
-    input_sub_mean = op.Sub(input, mean)
+    # Mean, var, and rstd computation and results are expected to be
+    # in higher precision when inputs are float16.
+    upcast_input = op.Cast(input, to=FLOAT.dtype)
+    mean = op.ReduceMean(upcast_input, axes)
+    input_sub_mean = op.Sub(upcast_input, mean)
     sqr = op.Mul(input_sub_mean, input_sub_mean)
     var = op.ReduceMean(sqr, axes, keepdims=False)
     rstd = op.Div(1.0, op.Sqrt(var + eps))
     # Get mean again with size = [1, C]
-    mean = op.ReduceMean(input, axes, keepdims=False)
+    mean = op.ReduceMean(upcast_input, axes, keepdims=False)
     # NOTE: Fixed to be FLOAT dtype
     running_mean = op.Cast(running_mean, to=FLOAT.dtype)
     running_var = op.Cast(running_var, to=FLOAT.dtype)

{onnxscript-0.1.0.dev20240123 → onnxscript-0.1.0.dev20240125}/onnxscript/tests/function_libs/torch_lib/ops_test_data.py

@@ -1838,7 +1838,6 @@ TESTED_TORCHLIB_OPS: tuple[TorchLibOpInfo, ...] = (
         "native_batch_norm",
         core_ops.aten_native_batch_norm,
         trace_only=True,
-        tolerance={torch.float16: (9e-3, 7e-4)},
     ),
     TorchLibOpInfo(
         "ops.aten._native_batch_norm_legit", core_ops.aten_native_batch_norm, trace_only=True

{onnxscript-0.1.0.dev20240123 → onnxscript-0.1.0.dev20240125/onnxscript.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: onnxscript
-Version: 0.1.0.dev20240123
+Version: 0.1.0.dev20240125
 Summary: Naturally author ONNX functions and models using a subset of Python
 Home-page: https://onnxscript.ai/
 Author-email: Microsoft Corporation <onnx@microsoft.com>