diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000000..a7555ac2a6 --- /dev/null +++ b/THIRD_PARTY_NOTICES.md @@ -0,0 +1,42 @@ +# Third-party notices + +## SciPy XSF Cephes `ndtr.h` coefficients + +`onnxscript/function_libs/torch_lib/ops/special.py` adapts the `ndtr_P`, +`ndtr_Q`, `ndtr_R`, `ndtr_S`, `ndtr_T`, and `ndtr_U` coefficient tables from +[SciPy XSF revision 5dbdff8de0dab99b475076612ea227d3f29d6cf6](https://github.com/scipy/scipy/blob/5dbdff8de0dab99b475076612ea227d3f29d6cf6/subprojects/xsf/include/xsf/cephes/ndtr.h). + +The source file identifies the original Cephes Math Library Release 2.2 as +copyright 1984, 1987, 1988, 1992 by Stephen L. Moshier. SciPy's 2024 C++ +translation is distributed under the BSD 3-Clause License: + +```text +BSD 3-Clause License + +Copyright (c) 2024, SciPy + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +``` diff --git a/onnxscript/function_libs/torch_lib/ops/special.py b/onnxscript/function_libs/torch_lib/ops/special.py index 1b123394d3..2ac9b57b1a 100644 --- a/onnxscript/function_libs/torch_lib/ops/special.py +++ b/onnxscript/function_libs/torch_lib/ops/special.py @@ -15,6 +15,8 @@ import math from typing import Optional, Sequence +from onnxscript import ir +from onnxscript.function_libs.torch_lib.ops import common as common_ops from onnxscript.function_libs.torch_lib.registration import torch_op from onnxscript.function_libs.torch_lib.tensor_typing import TFloat from onnxscript.onnx_opset import opset18 as op @@ -22,6 +24,78 @@ _MATH_PI = math.pi +# Coefficients adapted from SciPy XSF's Cephes ``ndtr.h`` (pinned source revision +# 5dbdff8de0dab99b475076612ea227d3f29d6cf6). The original Cephes Math Library +# Release 2.2 is copyright Stephen L. Moshier (1984, 1987, 1988, 1992); SciPy's +# C++ translation is BSD-3-Clause. See THIRD_PARTY_NOTICES.md for the license. +_ERFCX_P = ( + 2.46196981473530512524e-10, + 5.64189564831068821977e-1, + 7.46321056442269912687, + 4.86371970985681366614e1, + 1.96520832956077098242e2, + 5.26445194995477358631e2, + 9.34528527171957607540e2, + 1.02755188689515710272e3, + 5.57535335369399327526e2, +) +_ERFCX_Q = ( + 1.0, + 1.32281951154744992508e1, + 8.67072140885989742329e1, + 3.54937778887819891062e2, + 9.75708501743205489753e2, + 1.82390916687909736289e3, + 2.24633760818710981792e3, + 1.65666309194161350182e3, + 5.57535340817727675546e2, +) +_ERFCX_R = ( + 5.64189583547755073984e-1, + 1.27536670759978104416, + 5.01905042251180477414, + 6.16021097993053585195, + 7.40974269950448939160, + 2.97886665372100240670, +) +_ERFCX_S = ( + 1.0, + 2.26052863220117276590, + 9.39603524938001434673, + 1.20489539808096656605e1, + 1.70814450747565897222e1, + 9.60896809063285878198, + 3.36907645100081516050, +) +_ERFCX_T = ( + 9.60497373987051638749, + 9.00260197203842689217e1, + 2.23200534594684319226e3, + 7.00332514112805075473e3, + 5.55923013010394962768e4, +) +_ERFCX_U = ( + 1.0, + 3.35617141647503099647e1, + 5.21357949780152679795e2, + 4.59432382970980127987e3, + 2.26290000613890934246e4, + 4.92673942608635921086e4, +) + + +def _erfcx_constant(value: float, like: TFloat) -> TFloat: + """Creates a coefficient with float64 source precision and the input dtype.""" + return op.CastLike(common_ops.constant(value, dtype=ir.DataType.DOUBLE), like) + + +def _erfcx_polynomial(coefficients: Sequence[float], x: TFloat) -> TFloat: + """Emits Horner evaluation; the fixed loop is unrolled while tracing.""" + result = _erfcx_constant(coefficients[0], x) + for coefficient in coefficients[1:]: + result = result * x + _erfcx_constant(coefficient, x) + return result + def aten_special_airy_ai(x: TensorType) -> TensorType: """special_airy_ai(Tensor x) -> Tensor""" @@ -103,11 +177,57 @@ def aten_special_erfc(self: TFloat) -> TFloat: return op.Sub(1, op.Erf(self)) -@torch_op("aten::special_erfcx") +def _aten_special_erfcx(self: TFloat) -> TFloat: + """special_erfcx(Tensor self) -> Tensor""" + + # erfcx(x) is evaluated as a positive function of |x|, then reflected for + # negative x. Bound each rational approximation's input before evaluating it + # because ONNX Where evaluates both branches. + abs_self = op.Abs(self) + zero = _erfcx_constant(0.0, self) + one = _erfcx_constant(1.0, self) + two = _erfcx_constant(2.0, self) + eight = _erfcx_constant(8.0, self) + + central_x = op.Where(op.Less(abs_self, one), abs_self, zero) + central_z = central_x * central_x + central_p = _erfcx_polynomial(_ERFCX_T, central_z) + central_q = _erfcx_polynomial(_ERFCX_U, central_z) + central_erf = central_x * central_p / central_q + central = op.Exp(central_z) * (one - central_erf) + + middle_mask = op.And(op.GreaterOrEqual(abs_self, one), op.Less(abs_self, eight)) + middle_x = op.Where(middle_mask, abs_self, one) + middle_p = _erfcx_polynomial(_ERFCX_P, middle_x) + middle_q = _erfcx_polynomial(_ERFCX_Q, middle_x) + middle = middle_p / middle_q + + tail_x = op.Where(op.GreaterOrEqual(abs_self, eight), abs_self, eight) + tail_r = op.Div(one, tail_x) + # The tail's denominator has one higher degree than its numerator. Reversing + # the polynomials in 1 / |x| avoids overflow for large finite inputs. + tail_p = _erfcx_polynomial(_ERFCX_R[::-1], tail_r) + tail_q = _erfcx_polynomial(_ERFCX_S[::-1], tail_r) + tail = tail_r * tail_p / tail_q + + positive = op.Where( + op.Less(abs_self, one), central, op.Where(op.Less(abs_self, eight), middle, tail) + ) + reflected = op.Sub(op.Mul(two, op.Exp(op.Mul(self, self))), positive) + result = op.Where(op.Less(self, zero), reflected, positive) + return op.Where(op.IsNaN(self), self, result) + + +@torch_op("aten::special_erfcx", trace_only=True) def aten_special_erfcx(self: TFloat) -> TFloat: """special_erfcx(Tensor self) -> Tensor""" - return op.Mul(op.Exp(op.Pow(self, 2)), op.Sub(1, op.Erf(self))) + # The degree-eight middle polynomial overflows float16 even though the + # final ratio is finite. Evaluate low-precision inputs in float32, then + # restore the requested dtype. Float32 and float64 retain their precision. + if self.dtype in (ir.DataType.FLOAT16, ir.DataType.BFLOAT16): + return op.CastLike(_aten_special_erfcx(op.Cast(self, to=ir.DataType.FLOAT)), self) + return _aten_special_erfcx(self) def aten_special_erfinv(self: TensorType) -> TensorType: diff --git a/pyproject.toml b/pyproject.toml index 2e0d826460..cfb81362c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -34,6 +34,9 @@ dependencies = [ "typing_extensions>=4.10", ] +[tool.setuptools] +license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md"] + [tool.setuptools.packages.find] include = ["onnxscript*"] diff --git a/tests/function_libs/torch_lib/ops_test_data.py b/tests/function_libs/torch_lib/ops_test_data.py index 9c1f9088ac..77cf35ca0c 100644 --- a/tests/function_libs/torch_lib/ops_test_data.py +++ b/tests/function_libs/torch_lib/ops_test_data.py @@ -709,9 +709,7 @@ def _where_input_wrangler( TorchLibOpInfo( "expm1", special_ops.aten_special_expm1, tolerance={torch.float16: (1e-2, 2e-4)} ), - TorchLibOpInfo("special.erfcx", special_ops.aten_special_erfcx).xfail( - reason="fixme: The implementation is numerically unstable: https://github.com/microsoft/onnxscript/issues/1223" - ), + TorchLibOpInfo("special.erfcx", special_ops.aten_special_erfcx), TorchLibOpInfo( "ops.aten.fake_quantize_per_channel_affine", core_ops.aten_fake_quantize_per_channel_affine, diff --git a/tests/function_libs/torch_lib/special_test.py b/tests/function_libs/torch_lib/special_test.py new file mode 100644 index 0000000000..983627a8ea --- /dev/null +++ b/tests/function_libs/torch_lib/special_test.py @@ -0,0 +1,114 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +"""Focused TorchLib tests for special functions.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest +import torch + +from onnxscript.function_libs.torch_lib.ops import special +from tests.function_libs.torch_lib import ops_test_common + + +@pytest.mark.parametrize( + ("dtype", "rtol", "atol"), + ((np.float16, 2e-3, 0), (np.float32, 1.3e-6, 0), (np.float64, 2e-14, 0)), +) +def test_erfcx_ort_matches_torch_at_boundaries_and_special_values(dtype, rtol, atol): + """Exercises each approximation range and the low-precision working path.""" + + values = np.array( + [ + -np.inf, + -30.0, + -26.0, + -12.0, + -8.001, + -8.0, + -7.999, + -1.001, + -1.0, + -0.999, + -0.0, + 0.0, + 0.999, + 1.0, + 1.001, + 7.999, + 8.0, + 8.001, + 12.0, + 30.0, + np.inf, + np.nan, + ], + dtype=dtype, + ) + with np.errstate(over="ignore"): + expected = torch.special.erfcx(torch.from_numpy(values).float()).numpy().astype(dtype) + if dtype is np.float64: + expected = torch.special.erfcx(torch.from_numpy(values)).numpy() + actual = ops_test_common.graph_executor("test_erfcx", [torch.from_numpy(expected)])( + special.aten_special_erfcx, (values,), {} + )[0] + + np.testing.assert_allclose(actual, expected, rtol=rtol, atol=atol, equal_nan=True) + assert np.isposinf(actual[0]) + assert actual[-2] == 0 + assert np.isnan(actual[-1]) + + +@pytest.mark.parametrize("dtype", (np.float32, np.float64)) +def test_erfcx_ort_has_correct_large_positive_asymptote(dtype): + """Checks the reciprocal tail form that prevents polynomial overflow.""" + + values = np.array([8.0, 12.0, 30.0, np.finfo(dtype).max], dtype=dtype) + expected = torch.special.erfcx(torch.from_numpy(values)).numpy() + actual = ops_test_common.graph_executor("test_erfcx_tail", [torch.from_numpy(expected)])( + special.aten_special_erfcx, (values,), {} + )[0] + + rtol = 1.3e-6 if dtype is np.float32 else 2e-14 + np.testing.assert_allclose(actual, expected, rtol=rtol, atol=0) + np.testing.assert_allclose(values[-1] * actual[-1], 1 / math.sqrt(math.pi), rtol=rtol) + + +@pytest.mark.parametrize("dtype", (np.float16, np.float32, np.float64)) +@pytest.mark.parametrize("shape", ((), (0,), (2, 0, 3), (2, 3))) +def test_erfcx_ort_preserves_shape_and_dtype(dtype, shape): + values = np.ones(shape, dtype=dtype) + expected = torch.special.erfcx(torch.from_numpy(values).double()).numpy().astype(dtype) + actual = ops_test_common.graph_executor("test_erfcx_shape", [torch.from_numpy(expected)])( + special.aten_special_erfcx, (values,), {} + )[0] + assert actual.shape == values.shape + assert actual.dtype == values.dtype + np.testing.assert_allclose(actual, expected, rtol=2e-3 if dtype is np.float16 else 1e-6) + + +@pytest.mark.parametrize("dtype", (np.float32, np.float64)) +def test_erfcx_ort_across_approximation_intervals(dtype): + # Include adjacent representable values at each piecewise boundary. + boundaries = np.array([-8, -1, 0, 1, 8], dtype=dtype) + values = np.concatenate( + [ + np.linspace(-9, 0, 257, dtype=dtype), + np.linspace(0, 1, 257, dtype=dtype), + np.linspace(1, 8, 257, dtype=dtype), + np.geomspace(8, 1e30 if dtype is np.float32 else 1e300, 257).astype(dtype), + boundaries, + np.nextafter(boundaries, -np.inf), + np.nextafter(boundaries, np.inf), + ] + ) + expected = torch.special.erfcx(torch.from_numpy(values)).numpy() + actual = ops_test_common.graph_executor( + "test_erfcx_intervals", [torch.from_numpy(expected)] + )(special.aten_special_erfcx, (values,), {})[0] + np.testing.assert_allclose( + actual, expected, rtol=1.3e-6 if dtype is np.float32 else 2e-14, atol=0 + )