revisions

Author: lucascolley

src/array_api_extra/testing.py

@@ -31,7 +31,13 @@
 from ._lib._utils._helpers import jax_autojit, pickle_flatten, pickle_unflatten
 from ._lib._utils._typing import Array, Device
 
-__all__ = ["lazy_xp_function", "patch_lazy_xp_functions"]
+__all__ = [
+    "assert_close",
+    "assert_equal",
+    "assert_less",
+    "lazy_xp_function",
+    "patch_lazy_xp_functions",
+]
 
 if TYPE_CHECKING:  # pragma: no cover
     # TODO import override from typing (requires Python >=3.12)
@@ -467,7 +473,7 @@ def revert_on_exit() -> Generator[None]:  # numpydoc ignore=GL08
     return revert_on_exit()
 
 
-class CountingDaskScheduler(SchedulerGetCallable):
+class _CountingDaskScheduler(SchedulerGetCallable):
     """
     Dask scheduler that counts how many times `dask.compute` is called.
 
@@ -527,7 +533,7 @@ def _dask_wrap(
 
     @wraps(func)
     def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
-        scheduler = CountingDaskScheduler(n, msg)
+        scheduler = _CountingDaskScheduler(n, msg)
         with dask.config.set({"scheduler": scheduler}):  # pyright: ignore[reportPrivateImportUsage]
             out = func(*args, **kwargs)
 
@@ -541,7 +547,7 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
     return wrapper
 
 
-def require_numpy() -> ModuleType:  # numpydoc ignore=RT01
+def _require_numpy() -> ModuleType:  # numpydoc ignore=RT01
     """
     Import and return `numpy` if it is available, otherwise raise informative error.
     """
@@ -588,7 +594,7 @@ def _check_ns_shape_dtype(
     -------
     Actual array, desired array, their array namespace, the numpy module.
     """
-    np = require_numpy()
+    np = _require_numpy()
 
     actual_xp = array_namespace(actual)  # Raises on Python scalars and lists
     desired_xp = array_namespace(desired)
@@ -631,7 +637,7 @@ def _check_ns_shape_dtype(
             desired_shape = cast(tuple[float, ...], desired.shape)
 
     if check_shape:
-        msg = f"shapes do not match: {actual_shape} != f{desired_shape}"
+        msg = f"shapes do not match: {actual_shape} != {desired_shape}"
         assert actual_shape == desired_shape, msg
     elif desired.ndim > 0:
         # Ignore shape, but check flattened size. This is normally done by
@@ -640,7 +646,7 @@ def _check_ns_shape_dtype(
         # This check excludes 0d arrays as they are special-cased in NumPy.
         actual_size = math.prod(actual_shape)
         desired_size = math.prod(desired_shape)
-        msg = f"sizes do not match: {actual_size} != f{desired_size}"
+        msg = f"sizes do not match: {actual_size} != {desired_size}"
         assert actual_size == desired_size, msg
 
     if check_dtype:
@@ -665,7 +671,7 @@ def _as_numpy_array(  # numpydoc ignore=PR01,RT01
     """
     Convert array to NumPy, bypassing GPU-CPU transfer guards and densification guards.
     """
-    np = require_numpy()
+    np = _require_numpy()
     if is_cupy_namespace(xp):
         return xp.asnumpy(array)
     if is_pydata_sparse_namespace(xp):
@@ -708,7 +714,7 @@ def assert_close(
     xp: ModuleType | None = None,
 ) -> None:
     """
-    Check that two arrays are close, up to a tolerance.
+    Check that two arrays are close, up to tolerance ``atol + rtol * abs(desired)``.
 
     This is an interface to :func:`numpy.testing.assert_allclose` which accepts
     any standard-compatible array and performs additional array namespace,
@@ -755,9 +761,12 @@ def assert_close(
     -----
     The default `atol` and `rtol` differ from ``xp.all(xpx.isclose(a, b))``.
     For inexact dtypes, the default `rtol` is
-    ``xp.finfo(actual.dtype).eps ** 0.5 * 4``, which is roughly halfway between
-    :math:`\\sqrt{\\epsilon}` and the default for :func:`numpy.testing.assert_allclose`,
-    ``1e-7``. For other dtypes, the default ``1e-7`` is used.
+    ``xp.finfo(actual.dtype).eps ** 0.5 * 4``, which for ``float64`` is roughly halfway
+    between :math:`\\sqrt{\\epsilon}` and the default for
+    :func:`numpy.testing.assert_allclose`, ``1e-7``.
+    This gives a more reasonable default for lower precision dtypes,
+    for example approximately ``1e-3`` for ``float32``.
+    For exact dtypes, the default ``1e-7`` is used.
 
     Array arguments to `atol` and `rtol` must be valid input to :class:`float`.
     """
@@ -880,14 +889,14 @@ def assert_less(
     verbose : bool, default: True
         Whether to include the conflicting arrays in the error message on failure.
     check_dtype : bool, default: True
-        Whether to check agreement between actual and desired dtypes.
+        Whether to check agreement between the dtypes of `x` and `y`.
     check_shape : bool, default: True
-        Whether to check agreement between actual and desired shapes.
+        Whether to check agreement between the shapes of `x` and `y`.
     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.
+        A standard-compatible namespace which `x` and `y` must match.
 
     Raises
     ------