ENH: testing: add assert_close_nulp

docs/api-testing.md

@@ -11,6 +11,7 @@ These testing utilities are intended to help transition from using
     :toctree: generated
 
     assert_close
+    assert_close_nulp
     assert_equal
     assert_less
 ```

src/array_api_extra/testing.py

@@ -33,6 +33,7 @@
 
 __all__ = [
     "assert_close",
+    "assert_close_nulp",
     "assert_equal",
     "assert_less",
     "lazy_xp_function",
@@ -748,6 +749,9 @@ def assert_close(
 
     Raises
     ------
+    AssertionError
+        If `actual` and `desired` are not equal up to the defined tolerance.
+
     ImportError
         If :mod:`numpy` is not importable in the Python environment.
 
@@ -842,6 +846,9 @@ def assert_equal(
 
     Raises
     ------
+    AssertionError
+        If `actual` and `desired` are not equal.
+
     ImportError
         If :mod:`numpy` is not importable in the Python environment.
 
@@ -900,6 +907,9 @@ def assert_less(
 
     Raises
     ------
+    AssertionError
+        If `x` is not strictly smaller than `y`, elementwise.
+
     ImportError
         If :mod:`numpy` is not importable in the Python environment.
 
@@ -916,3 +926,76 @@ def assert_less(
     x_np = _as_numpy_array(x, xp=xp)
     y_np = _as_numpy_array(y, xp=xp)
     np.testing.assert_array_less(x_np, y_np, err_msg=err_msg, verbose=verbose)
+
+
+def assert_close_nulp(
+    actual: Array,
+    desired: Array,
+    *,
+    nulp: int = 1,
+    check_dtype: bool = True,
+    check_shape: bool = True,
+    check_scalar: bool = False,
+    xp: ModuleType | None = None,
+) -> None:
+    """
+    Compare two arrays relatively to their spacing.
+
+    This is an interface to :func:`numpy.testing.assert_array_almost_equal_nulp`
+    which accepts any standard-compatible array and performs
+    additional array namespace, shape, and dtype checks.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+    nulp : int, optional
+        The maximum number of units in the last place
+        for the tolerance check. Default: ``1``.
+    check_dtype : bool, default: True
+        Whether to check agreement between actual and desired dtypes.
+    check_shape : bool, default: True
+        Whether to check agreement between actual and desired shapes.
+    check_scalar : bool, default: False
+        NumPy only: whether to check agreement between actual and desired types —
+        0-D :class:`numpy.ndarray` vs scalar (e.g. :class:`numpy.double`).
+    xp : array_namespace, optional
+        A standard-compatible namespace which `actual` and `desired` must match.
+
+    Raises
+    ------
+    AssertionError
+        If the spacing between `actual` and `desired` for one or more elements is \
+        larger than `nulp`.
+
+    ImportError
+        If :mod:`numpy` is not importable in the Python environment.
+
+    See Also
+    --------
+    assert_close : Similar function for inexact equality checks.
+    numpy.spacing : Spacing calculation for NumPy arrays.
+    numpy.testing.assert_array_almost_equal_nulp : Similar function for NumPy arrays.
+
+    Notes
+    -----
+    This is a relatively robust method to compare two arrays whose amplitude is
+    variable.
+
+    An assertion is raised if the following condition is not met::
+
+        abs(actual - desired) <= nulp * spacing(maximum(abs(actual), abs(desired)))
+
+    where ``spacing(x)`` is the distance between ``x`` and the nearest adjacent number
+    representable by in the data type of ``x``.
+    """
+    actual, desired, xp, np = _check_ns_shape_dtype(
+        actual, desired, check_dtype, check_shape, check_scalar, xp
+    )
+    if not _is_materializable(actual):
+        return
+    actual_np = _as_numpy_array(actual, xp=xp)
+    desired_np = _as_numpy_array(desired, xp=xp)
+    np.testing.assert_array_almost_equal_nulp(actual_np, desired_np, nulp=nulp)

tests/test_testing.py

@@ -17,6 +17,7 @@
 from array_api_extra.testing import (
     _as_numpy_array,
     assert_close,
+    assert_close_nulp,
     assert_equal,
     assert_less,
     lazy_xp_function,
@@ -40,6 +41,9 @@ def test_device(self, xp: ModuleType, device: Device):
 
 
 class TestAssertEqualCloseLess:
+    def np_err_msg(self, func: object):
+        return "not equal" if func is assert_close_nulp else "Mismatched elements"
+
     @pytest.mark.parametrize("func", [assert_equal, assert_close])
     def test_assert_equal_close_basic(self, xp: ModuleType, func: Callable[..., None]):
         func(xp.asarray(0), xp.asarray(0))
@@ -51,7 +55,9 @@ def test_assert_equal_close_basic(self, xp: ModuleType, func: Callable[..., None
         with pytest.raises(AssertionError, match="hello"):
             func(xp.asarray([1, 2]), xp.asarray([2, 1]), err_msg="hello")
 
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     def test_shape_dtype(self, xp: ModuleType, func: Callable[..., None]):
         with pytest.raises(AssertionError, match="shapes do not match"):
             func(xp.asarray([0]), xp.asarray([[0]]))
@@ -63,7 +69,9 @@ def test_shape_dtype(self, xp: ModuleType, func: Callable[..., None]):
     @pytest.mark.skip_xp_backend(
         Backend.NUMPY_READONLY, reason="test other ns vs. numpy"
     )
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     def test_namespace(self, xp: ModuleType, func: Callable[..., None]):
         with pytest.raises(
             AssertionError, match="Namespaces of actual and desired arrays do not match"
@@ -81,7 +89,9 @@ def test_namespace(self, xp: ModuleType, func: Callable[..., None]):
         ):
             func(xp.asarray(0), xp.asarray(0), xp=np)
 
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     def test_check_shape(self, xp: ModuleType, func: Callable[..., None]):
         a = xp.asarray([1] if func is assert_less else [2])
         b = xp.asarray(2)
@@ -91,12 +101,14 @@ def test_check_shape(self, xp: ModuleType, func: Callable[..., None]):
         with pytest.raises(AssertionError, match="shapes do not match"):
             func(a, b)
         func(a, b, check_shape=False)
-        with pytest.raises(AssertionError, match="Mismatched elements"):
+        with pytest.raises(AssertionError, match=self.np_err_msg(func)):
             func(a, c, check_shape=False)
         with pytest.raises(AssertionError, match="sizes do not match"):
             func(a, d, check_shape=False)
 
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     def test_check_dtype(self, xp: ModuleType, func: Callable[..., None]):
         a = xp.asarray(1 if func is assert_less else 2)
         b = xp.asarray(2, dtype=xp.int16)
@@ -105,10 +117,12 @@ def test_check_dtype(self, xp: ModuleType, func: Callable[..., None]):
         with pytest.raises(AssertionError, match="dtypes do not match"):
             func(a, b)
         func(a, b, check_dtype=False)
-        with pytest.raises(AssertionError, match="Mismatched elements"):
+        with pytest.raises(AssertionError, match=self.np_err_msg(func)):
             func(a, c, check_dtype=False)
 
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     @pytest.mark.xfail_xp_backend(
         Backend.SPARSE, reason="sparse [()] returns np.generic"
     )
@@ -125,7 +139,7 @@ def test_check_scalar(
                 func(a, b, check_scalar=True)
         else:
             func(a, b, check_scalar=True)
-        with pytest.raises(AssertionError, match="Mismatched elements"):
+        with pytest.raises(AssertionError, match=self.np_err_msg(func)):
             func(a, c, check_scalar=True)
 
     @pytest.mark.parametrize("dtype", ["int64", "float64"])
@@ -150,7 +164,9 @@ def test_assert_less(self, xp: ModuleType):
         with pytest.raises(AssertionError, match="Mismatched elements"):
             assert_less(xp.asarray([1, 1]), xp.asarray([2, 1]))
 
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     @pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array")
     @pytest.mark.skip_xp_backend(Backend.ARRAY_API_STRICTEST, reason="boolean indexing")
     def test_none_shape(self, xp: ModuleType, func: Callable[..., None]):
@@ -166,7 +182,7 @@ def test_none_shape(self, xp: ModuleType, func: Callable[..., None]):
             func(a, xp.asarray(2))
         with pytest.raises(AssertionError, match="shapes do not match"):
             func(a, xp.asarray([2, 3]))
-        with pytest.raises(AssertionError, match="Mismatched elements"):
+        with pytest.raises(AssertionError, match=self.np_err_msg(func)):
             func(a, xp.asarray([0]))
 
         # desired has shape=(None, )
@@ -178,10 +194,12 @@ def test_none_shape(self, xp: ModuleType, func: Callable[..., None]):
             func(xp.asarray(2), a)
         with pytest.raises(AssertionError, match="shapes do not match"):
             func(xp.asarray([2, 3]), a)
-        with pytest.raises(AssertionError, match="Mismatched elements"):
+        with pytest.raises(AssertionError, match=self.np_err_msg(func)):
             func(xp.asarray([4]), a)
 
-    @pytest.mark.parametrize("func", [assert_equal, assert_close, assert_less])
+    @pytest.mark.parametrize(
+        "func", [assert_equal, assert_close, assert_less, assert_close_nulp]
+    )
     def test_device(self, xp: ModuleType, device: Device, func: Callable[..., None]):
         a = xp.asarray([1] if func is assert_less else [2], device=device)
         b = xp.asarray([2], device=device)
@@ -195,6 +213,31 @@ def test_device(self, xp: ModuleType, device: Device, func: Callable[..., None])
         with pytest.raises(AssertionError, match="sizes do not match"):
             func(a, c, check_shape=False)
 
+    def test_assert_close_nulp(self, xp: ModuleType):
+        a = xp.asarray([1.0, 1e-10])
+        b = xp.asarray([1.0 + 2**-52, 1e-10])
+
+        assert_close_nulp(a, a, nulp=0)
+
+        assert_close_nulp(a, b, nulp=2)
+        assert_close_nulp(a, b, nulp=1)
+        with pytest.raises(AssertionError, match="not equal to 0 ULP"):
+            assert_close_nulp(a, b, nulp=0)
+
+        c = xp.asarray([1.0 + 2 * 2**-52, 1e-10])
+
+        assert_close_nulp(a, c, nulp=3)
+        assert_close_nulp(a, c, nulp=2)
+        with pytest.raises(AssertionError, match="not equal to 1 ULP"):
+            assert_close_nulp(a, c)
+
+        d = xp.asarray([1.0, 1e-10 + 5 * np.spacing(1e-10)])
+        assert_close_nulp(a, d, nulp=6)
+        assert_close_nulp(a, d, nulp=5)
+
+        with pytest.raises(AssertionError, match="not equal to 4 ULP"):
+            assert_close_nulp(a, d, nulp=4)
+
 
 def good_lazy(x: Array) -> Array:
     """A function that behaves well in Dask and jax.jit"""