ENH: new functions isclose and allclose

Author: crusaderky

Diff for: docs/api-reference.md

@@ -6,11 +6,13 @@
     :nosignatures:
     :toctree: generated
 
+    allclose
     at
     atleast_nd
     cov
     create_diagonal
     expand_dims
+    isclose
     kron
     nunique
     pad

Diff for: pixi.lock

@@ -293,6 +293,7 @@ checks = [
     "all",  # report on all checks, except the below
     "EX01", # most docstrings do not need an example
     "SA01", # data-apis/array-api-extra#87
+    "SA04", # Missing description for See Also cross-reference
     "ES01", # most docstrings do not need an extended summary
 ]
 exclude = [ # don't report on objects that match any of these regex

Diff for: pyproject.toml

@@ -1,6 +1,6 @@
 """Extra array functions built on top of the array API standard."""
 
-from ._delegation import pad
+from ._delegation import allclose, isclose, pad
 from ._lib._at import at
 from ._lib._funcs import (
     atleast_nd,
@@ -18,11 +18,13 @@
 # pylint: disable=duplicate-code
 __all__ = [
     "__version__",
+    "allclose",
     "at",
     "atleast_nd",
     "cov",
     "create_diagonal",
     "expand_dims",
+    "isclose",
     "kron",
     "nunique",
     "pad",

Diff for: src/array_api_extra/__init__.py

@@ -7,7 +7,7 @@
 from ._lib._utils._compat import array_namespace
 from ._lib._utils._typing import Array
 
-__all__ = ["pad"]
+__all__ = ["allclose", "isclose", "pad"]
 
 
 def _delegate(xp: ModuleType, *backends: Backend) -> bool:
@@ -29,6 +29,144 @@ def _delegate(xp: ModuleType, *backends: Backend) -> bool:
     return any(backend.is_namespace(xp) for backend in backends)
 
 
+def allclose(
+    a: Array,
+    b: Array,
+    *,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Return True if two arrays are element-wise equal within a tolerance.
+
+    This is a simple convenience reduction around `isclose`.
+
+    Parameters
+    ----------
+    a, b : Array
+        Input arrays to compare.
+    rtol : array_like, optional
+        The relative tolerance parameter.
+    atol : array_like, optional
+        The absolute tolerance parameter.
+    equal_nan : bool, optional
+        Whether to compare NaN's as equal. If True, NaN's in `a` will be considered
+        equal to NaN's in `b` in the output array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a` and `b`. Default: infer.
+
+    Returns
+    -------
+    Array
+        A 0-dimensional boolean array, containing `True` if all `a` is elementwise close
+        to `b` and `False` otherwise.
+
+    See Also
+    --------
+    isclose
+    math.isclose
+
+    Notes
+    -----
+    If `xp` is a lazy backend (e.g. Dask, JAX), you may not be able to test the result
+    contents with ``bool(allclose(a, b))`` or ``if allclose(a, b): ...``.
+    """
+    xp = array_namespace(a, b) if xp is None else xp
+    return xp.all(isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan, xp=xp))
+
+
+def isclose(
+    a: Array,
+    b: Array,
+    *,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Return a boolean array where two arrays are element-wise equal within a tolerance.
+
+    The tolerance values are positive, typically very small numbers. The relative
+    difference `(rtol * abs(b))` and the absolute difference atol are added together to
+    compare against the absolute difference between a and b.
+
+    NaNs are treated as equal if they are in the same place and if equal_nan=True. Infs
+    are treated as equal if they are in the same place and of the same sign in both
+    arrays.
+
+    Parameters
+    ----------
+    a, b : Array
+        Input arrays to compare.
+    rtol : array_like, optional
+        The relative tolerance parameter (see Notes).
+    atol : array_like, optional
+        The absolute tolerance parameter (see Notes).
+    equal_nan : bool, optional
+        Whether to compare NaN's as equal. If True, NaN's in `a` will be considered
+        equal to NaN's in `b` in the output array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a` and `b`. Default: infer.
+
+    Returns
+    -------
+    Array
+        A boolean array of shape broadcasted from `a` and `b`, containing `True` where
+        ``a`` is close to ``b``, and `False` otherwise.
+
+    Warnings
+    --------
+    The default atol is not appropriate for comparing numbers with magnitudes much
+    smaller than one ) (see notes).
+
+    See Also
+    --------
+    allclose
+    math.isclose
+
+    Notes
+    -----
+    For finite values, `isclose` uses the following equation to test whether two
+    floating point values are equivalent::
+
+        absolute(a - b) <= (atol + rtol * absolute(b))
+
+    Unlike the built-in `math.isclose`, the above equation is not symmetric in a and b,
+    so that `isclose(a, b)` might be different from `isclose(b, a)` in some rare
+    cases.
+
+    The default value of `atol` is not appropriate when the reference value `b` has
+    magnitude smaller than one. For example, it is unlikely that ``a = 1e-9`` and
+    ``b = 2e-9`` should be considered "close", yet ``isclose(1e-9, 2e-9)`` is `True`
+    with default settings. Be sure to select atol for the use case at hand, especially
+    for defining the threshold below which a non-zero value in `a` will be considered
+    "close" to a very small or zero value in `b`.
+
+    The comparison of `a` and `b` uses standard broadcasting, which means that `a` and
+    `b` need not have the same shape in order for `isclose(a, b)` to evaluate to
+    `True`.
+
+    `isclose` is not defined for non-numeric data types. `bool` is considered a numeric
+    data-type for this purpose.
+    """
+    xp = array_namespace(a, b) if xp is None else xp
+
+    if _delegate(
+        xp,
+        Backend.NUMPY,
+        Backend.CUPY,
+        Backend.DASK,
+        Backend.JAX,
+        Backend.TORCH,
+    ):
+        return xp.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+    return _funcs.isclose(a, b, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+
 def pad(
     x: Array,
     pad_width: int | tuple[int, int] | list[tuple[int, int]],

Diff for: src/array_api_extra/_delegation.py

@@ -304,6 +304,41 @@ def expand_dims(
     return a
 
 
+def isclose(
+    a: Array,
+    b: Array,
+    *,
+    rtol: float = 1e-05,
+    atol: float = 1e-08,
+    equal_nan: bool = False,
+    xp: ModuleType | None = None,
+) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
+    xp = array_namespace(a, b) if xp is None else xp
+
+    a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating"))
+    b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating"))
+    if a_inexact or b_inexact:
+        # FIXME: use scipy's lazywhere to suppress warnings on inf
+        out = xp.abs(a - b) <= (atol + rtol * xp.abs(b))
+        out = xp.where(xp.isinf(a) & xp.isinf(b), xp.sign(a) == xp.sign(b), out)
+        if equal_nan:
+            out = xp.where(xp.isnan(a) & xp.isnan(b), xp.asarray(True), out)
+        return out
+
+    if xp.isdtype(a.dtype, "bool") or xp.isdtype(b.dtype, "bool"):
+        if atol >= 1 or rtol >= 1:
+            return xp.ones_like(a == b)
+        return a == b
+
+    # integer types
+    atol = int(atol)
+    if rtol == 0:
+        return xp.abs(a - b) <= atol
+    nrtol = int(1.0 / rtol)
+    return xp.abs(a - b) <= (atol + xp.abs(b) // nrtol)
+
+
 def kron(a: Array, b: Array, /, *, xp: ModuleType | None = None) -> Array:
     """
     Kronecker product of two arrays.