From e79389a367d5c77770bac04b7ef011dc16d8c2c7 Mon Sep 17 00:00:00 2001 From: Neil Girdhar Date: Thu, 22 Oct 2026 02:41:18 -0400 Subject: [PATCH] ENH: add `one_minus_square` Compute `1 - x**2` as `(1 + x) * (1 - x)`, which avoids catastrophic cancellation when `|x|` is close to 1. On JAX >= 0.11.2, delegate to `jax.lax.one_minus_square`, whose derivative is defined as `-2x` and so stays accurate near zero. Co-Authored-By: Claude Opus 5.5 --- docs/sphinx/api-elementwise.md | 1 + src/array_api_extra/__init__.py | 11 ++- src/array_api_extra/_agnostic/_elementwise.py | 7 ++ src/array_api_extra/_elementwise.py | 67 ++++++++++++++++++- tests/main/test_elementwise.py | 41 ++++++++++++ 5 files changed, 125 insertions(+), 2 deletions(-) diff --git a/docs/sphinx/api-elementwise.md b/docs/sphinx/api-elementwise.md index f75022fe..98d0e6fc 100644 --- a/docs/sphinx/api-elementwise.md +++ b/docs/sphinx/api-elementwise.md @@ -11,6 +11,7 @@ deg2rad isclose nan_to_num + one_minus_square rad2deg sinc ``` diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 8727bb40..ab07defb 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -5,7 +5,15 @@ from ._agnostic._inspection import default_dtype from ._at import at from ._creation import create_diagonal, one_hot -from ._elementwise import angle, deg2rad, isclose, nan_to_num, rad2deg, sinc +from ._elementwise import ( + angle, + deg2rad, + isclose, + nan_to_num, + one_minus_square, + rad2deg, + sinc, +) from ._indexing import diag_indices, tril_indices, triu_indices, unravel_index from ._lazy import lazy_apply from ._linalg import kron @@ -42,6 +50,7 @@ "nansum", "nunique", "one_hot", + "one_minus_square", "pad", "partition", "rad2deg", diff --git a/src/array_api_extra/_agnostic/_elementwise.py b/src/array_api_extra/_agnostic/_elementwise.py index 5f41426f..eb111fe2 100644 --- a/src/array_api_extra/_agnostic/_elementwise.py +++ b/src/array_api_extra/_agnostic/_elementwise.py @@ -14,6 +14,7 @@ "deg2rad", "isclose", "nan_to_num", + "one_minus_square", "rad2deg", "sinc", ] @@ -315,3 +316,9 @@ def rad2deg(x: Array, /, *, xp: ArrayNamespace) -> Array: # numpydoc ignore=PR01,RT01 """See docstring in `array_api_extra._elementwise`.""" return x * 180 / xp.pi + + +def one_minus_square(x: Array, /) -> Array: + # numpydoc ignore=PR01,RT01 + """See docstring in `array_api_extra._elementwise`.""" + return (1 + x) * (1 - x) diff --git a/src/array_api_extra/_elementwise.py b/src/array_api_extra/_elementwise.py index 23b29a67..e32c2f9d 100644 --- a/src/array_api_extra/_elementwise.py +++ b/src/array_api_extra/_elementwise.py @@ -4,7 +4,15 @@ from ._lib import _compat, _helpers from ._lib._typing import Array, ArrayNamespace -__all__ = ["angle", "deg2rad", "isclose", "nan_to_num", "rad2deg", "sinc"] +__all__ = [ + "angle", + "deg2rad", + "isclose", + "nan_to_num", + "one_minus_square", + "rad2deg", + "sinc", +] def angle(z: Array, /, *, deg: bool = False, xp: ArrayNamespace | None = None) -> Array: @@ -340,6 +348,63 @@ def nan_to_num( return _agnostic._elementwise.nan_to_num(y, fill_value=fill_value, xp=xp) +def one_minus_square(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: + r""" + Compute :math:`1 - x^2` element-wise. + + The result is evaluated as :math:`(1 + x)(1 - x)`, which avoids the + catastrophic cancellation of the naive formula when :math:`|x|` is close to + 1. + + Parameters + ---------- + x : array + Input array. Must have a numeric dtype. + xp : array_namespace, optional + The standard-compatible namespace for `x`. Default: infer. + + Returns + ------- + array + ``1 - x**2`` calculated element-wise, with the same shape and dtype as `x`. + + See Also + -------- + jax.lax.one_minus_square : JAX equivalent, which additionally defines its + derivative as :math:`-2x` for accurate gradients near zero. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + >>> xpx.one_minus_square(xp.asarray([0.0, 0.5, 1.0, 2.0]), xp=xp) + Array([ 1. , 0.75, 0. , -3. ], dtype=array_api_strict.float64) + + Near 1, the result keeps full relative precision: + + >>> x = xp.asarray(1 - 2.0**-30) + >>> float(xpx.one_minus_square(x, xp=xp)) == 2.0**-29 - 2.0**-60 + True + >>> float(1 - x**2) == 2.0**-29 - 2.0**-60 + False + """ + if xp is None: + xp = _compat.array_namespace(x) + + if not xp.isdtype(x.dtype, "numeric"): + msg = "`x` must have a numeric dtype." + raise TypeError(msg) + + if _compat.is_jax_namespace(xp): + import jax + + # Added in JAX 0.11.2. + if hasattr(jax.lax, "one_minus_square"): + return jax.lax.one_minus_square(x) + + return _agnostic._elementwise.one_minus_square(x) + + def sinc(x: Array, /, *, xp: ArrayNamespace | None = None) -> Array: r""" Return the normalized sinc function. diff --git a/tests/main/test_elementwise.py b/tests/main/test_elementwise.py index 213dbd48..c13f548d 100644 --- a/tests/main/test_elementwise.py +++ b/tests/main/test_elementwise.py @@ -15,6 +15,7 @@ deg2rad, isclose, nan_to_num, + one_minus_square, rad2deg, sinc, ) @@ -28,6 +29,7 @@ lazy_xp_function(deg2rad) lazy_xp_function(isclose) lazy_xp_function(nan_to_num) +lazy_xp_function(one_minus_square) lazy_xp_function(rad2deg) lazy_xp_function(sinc) @@ -697,6 +699,45 @@ def test_device(self, xp: ArrayNamespace, device: Device): assert get_device(angle(a)) == device +class TestOneMinusSquare: + def test_basic(self, xp: ArrayNamespace): + x = xp.asarray([-2.0, -1.0, 0.0, 0.5, 1.0, 3.0]) + expected = xp.asarray([-3.0, 0.0, 1.0, 0.75, 0.0, -8.0]) + assert_equal(one_minus_square(x), expected) + + def test_near_one(self, xp: ArrayNamespace): + x = xp.asarray([1 - 2.0**-30, -(1 - 2.0**-30)], dtype=xp.float64) + expected = xp.asarray([2.0**-29 - 2.0**-60] * 2, dtype=xp.float64) + assert_equal(one_minus_square(x), expected) + + @pytest.mark.parametrize("dtype_name", ["float32", "float64"]) + def test_dtype(self, xp: ArrayNamespace, dtype_name: str): + dtype = getattr(xp, dtype_name) + x = xp.asarray([0.5], dtype=dtype) + actual = one_minus_square(x, xp=xp) + assert actual.dtype == dtype + assert_equal(actual, xp.asarray([0.75], dtype=dtype)) + + def test_integral(self, xp: ArrayNamespace): + x = xp.asarray([0, 1, 2, -3], dtype=xp.int32) + actual = one_minus_square(x, xp=xp) + assert_equal(actual, xp.asarray([1, 0, -3, -8], dtype=xp.int32)) + + def test_complex(self, xp: ArrayNamespace): + x = xp.asarray([0.5 + 1j], dtype=xp.complex128) + actual = one_minus_square(x, xp=xp) + assert_close(actual, xp.asarray([1.75 - 1j], dtype=xp.complex128)) + + def test_bool(self, xp: ArrayNamespace): + x = xp.asarray([True]) + with pytest.raises(TypeError, match="numeric dtype"): + _ = one_minus_square(x, xp=xp) + + def test_device(self, xp: ArrayNamespace, device: Device): + x = xp.asarray([0.5], device=device) + assert get_device(one_minus_square(x)) == device + + class TestDeg2Rad: def test_basic(self, xp: ArrayNamespace): x = xp.asarray([0.0, 90.0, 180.0, 270.0, 360.0])