Add tolerance limits

microsoft · Jul 16, 2024 · 36f96d5 · 36f96d5
1 parent 4ad38ad
commit 36f96d5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -4666,11 +4666,11 @@ def aten_le_bool(self: BOOL, other: BOOL) -> BOOL:
 def aten_lerp(self: TReal, end: TReal, weight: TReal) -> TReal:
     """lerp.Tensor(Tensor self, Tensor end, Tensor weight) -> Tensor"""
 
-    diff = op.Sub(end, self)
+    diff = op.CastLike(op.Sub(end, self), weight)
     return op.Where(
         op.Less(weight, 0.5),
-        op.Add(self, op.Mul(weight, diff)),
-        op.Sub(end, op.Mul(diff, op.Sub(1.0, weight)))
+        op.Add(self, op.CastLike(op.Mul(weight, diff), self)),
+        op.Sub(end, op.CastLike(op.Mul(diff, op.Sub(1.0, weight)), end))
     )
 
 

diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py
@@ -900,7 +900,11 @@ def _where_input_wrangler(
     TorchLibOpInfo("log", core_ops.aten_log),
     TorchLibOpInfo("le", core_ops.aten_le),
     TorchLibOpInfo("le_bool", core_ops.aten_le_bool),
-    TorchLibOpInfo("lerp", core_ops.aten_lerp),
+    TorchLibOpInfo(
+        "lerp",
+        core_ops.aten_lerp,
+        tolerance={torch.float16: (2e-3, 2e-1)},
+    ),
     TorchLibOpInfo("log10", core_ops.aten_log10),
     TorchLibOpInfo("log1p", core_ops.aten_log1p),
     TorchLibOpInfo(
@@ -1021,7 +1025,11 @@ def _where_input_wrangler(
     TorchLibOpInfo("mT", core_ops.aten_mT_complex, complex=True),
     TorchLibOpInfo("mul", core_ops.aten_mul),
     TorchLibOpInfo("mul", core_ops.aten_mul_complex, complex=True),
-    TorchLibOpInfo("mv", core_ops.aten_mv),
+    TorchLibOpInfo(
+        "mv",
+        core_ops.aten_mv,
+        tolerance={torch.float16: (3e-2, 1e-2)},
+    ),
     TorchLibOpInfo("narrow", core_ops.aten_narrow),
     TorchLibOpInfo("ops.aten.native_dropout", core_ops.aten_native_dropout),
     TorchLibOpInfo("ne", core_ops.aten_ne),