Repository: microsoft/onnxscript

Issue:
[torchlib] Implement `erfc`



Discussion:

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
@@ -2308,19 +2308,6 @@ def aten_equal(self: TTensor, other: TTensor) -> BOOL:
     return op.Equal(sum_of_abs, 0)
 
 
-@torch_op("aten::erf")
-def aten_erf(self: TReal) -> TReal:
-    """erf(Tensor self) -> Tensor"""
-
-    return op.Erf(self)
-
-
-def aten_erfc(self: TensorType) -> TensorType:
-    """erfc(Tensor self) -> Tensor"""
-
-    raise NotImplementedError()
-
-
 def aten_erfinv(self: TensorType) -> TensorType:
     """erfinv(Tensor self) -> Tensor"""
 
diff --git a/onnxscript/function_libs/torch_lib/ops/special.py b/onnxscript/function_libs/torch_lib/ops/special.py
--- a/onnxscript/function_libs/torch_lib/ops/special.py
+++ b/onnxscript/function_libs/torch_lib/ops/special.py
@@ -15,7 +15,7 @@
 
 from onnxscript import FLOAT
 from onnxscript.function_libs.torch_lib.registration import torch_op
-from onnxscript.function_libs.torch_lib.tensor_typing import TFloatOrBFloat16
+from onnxscript.function_libs.torch_lib.tensor_typing import TFloatOrBFloat16, TReal
 from onnxscript.onnx_opset import opset18 as op
 from onnxscript.onnx_types import TensorType
 
@@ -86,22 +86,25 @@ def aten_special_entr(self: TensorType) -> TensorType:
     raise NotImplementedError()
 
 
-def aten_special_erf(self: TensorType) -> TensorType:
-    """special_erf(Tensor self) -> Tensor"""
+@torch_op("aten::erf")
+def aten_special_erf(self: TReal) -> TReal:
+    """erf(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.Erf(self)
 
 
-def aten_special_erfc(self: TensorType) -> TensorType:
-    """special_erfc(Tensor self) -> Tensor"""
+@torch_op("aten::erfc")
+def aten_special_erfc(self: TReal) -> TReal:
+    """erfc(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.Sub(1, op.Erf(self))
 
 
-def aten_special_erfcx(self: TensorType) -> TensorType:
+@torch_op("aten::erfcx")
+def aten_special_erfcx(self: TFloatOrBFloat16) -> TFloatOrBFloat16:
     """special_erfcx(Tensor self) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.Mul(op.Exp(op.Pow(self, 2)), op.Erf(self))
 
 
 def aten_special_erfinv(self: TensorType) -> TensorType:


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
@@ -552,7 +552,11 @@ def _where_input_wrangler(
     TorchLibOpInfo("exp2", core_ops.aten_exp2),
     TorchLibOpInfo("expand", core_ops.aten_expand),
     TorchLibOpInfo("expand_as", core_ops.aten_expand_as),
-    TorchLibOpInfo("erf", core_ops.aten_erf),
+    TorchLibOpInfo("erf", special_ops.aten_special_erf),
+    TorchLibOpInfo(
+        "erfc", special_ops.aten_special_erfc, tolerance={torch.float16: (1e-2, 2e-4)}
+    ),
+    # TorchLibOpInfo("erfcx", special_ops.aten_special_erfcx),  # not in OPS_DB
     TorchLibOpInfo("fill", core_ops.aten_fill),
     TorchLibOpInfo("flip", core_ops.aten_flip, input_wrangler=_flip_input_wrangler),
     TorchLibOpInfo("floor", core_ops.aten_floor),