diff --git a/src/array_api_extra/testing/__init__.py b/src/array_api_extra/testing/__init__.py index a57a88a6..d57735ca 100644 --- a/src/array_api_extra/testing/__init__.py +++ b/src/array_api_extra/testing/__init__.py @@ -7,6 +7,8 @@ assert_less, lazy_xp_function, patch_lazy_xp_functions, + yield_namespace_device_dtype_combinations, + yield_namespaces, ) __all__ = [ @@ -16,4 +18,6 @@ "assert_less", "lazy_xp_function", "patch_lazy_xp_functions", + "yield_namespace_device_dtype_combinations", + "yield_namespaces", ] diff --git a/src/array_api_extra/testing/_testing.py b/src/array_api_extra/testing/_testing.py index 575866e9..9244a5bd 100644 --- a/src/array_api_extra/testing/_testing.py +++ b/src/array_api_extra/testing/_testing.py @@ -6,6 +6,7 @@ import enum import functools import inspect +import itertools import math import typing import warnings @@ -35,6 +36,8 @@ "assert_less", "lazy_xp_function", "patch_lazy_xp_functions", + "yield_namespace_device_dtype_combinations", + "yield_namespaces", ] @@ -1008,3 +1011,89 @@ def assert_close_nulp( 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) + + +def yield_namespaces(include_numpy_namespaces: bool = True) -> Generator[str]: + """Yield supported namespace. + + This is meant to be used for testing purposes only. + + Parameters + ---------- + include_numpy_namespaces : bool, default=True + If True, also yield numpy namespaces. + + Returns + ------- + str + The name of the Array API namespace. + """ + _NUMPY_NAMESPACE_NAMES = {"numpy", "sklearn.externals.array_api_compat.numpy"} + for array_namespace in [ + # The following is used to test the array_api_compat wrapper when + # array_api_dispatch is enabled: in particular, the arrays used in the + # tests are regular numpy arrays without any "device" attribute. + "numpy", + # Stricter NumPy-based Array API implementation. The + # array_api_strict.Array instances always have a dummy "device" attribute. + "array_api_strict", + "cupy", + "torch", + "dpnp", + ]: + if not include_numpy_namespaces and array_namespace in _NUMPY_NAMESPACE_NAMES: + continue + yield array_namespace + + +def yield_namespace_device_dtype_combinations( + include_numpy_namespaces: bool = True, +) -> Generator[Any]: + """Yield supported namespace, device_name, dtype_name tuples for testing. + + Use this to test that an estimator works with all combinations. + Pass the yielded values to `_array_api_for_tests` which returns (xp, device) + for array allocation. + + Parameters + ---------- + include_numpy_namespaces : bool, default=True + If True, also yield numpy namespaces. + + Returns + ------- + array_namespace : str + The name of the Array API namespace. + + device_name : str or None + The name of the device on which to allocate the arrays. Can be None to + indicate that the default value should be used. + + dtype_name : str + The name of the data type to use for arrays. Can be None to indicate + that the default value should be used. + """ + for array_namespace in yield_namespaces( + include_numpy_namespaces=include_numpy_namespaces + ): + if array_namespace == "torch": + for device_name, dtype in itertools.product( + ("cpu", "cuda", "xpu"), ("float64", "float32") + ): + yield array_namespace, device_name, dtype + yield array_namespace, "mps", "float32" + + elif array_namespace == "dpnp": # pragma: nocover + # XXX: add "accelerator" device type? + for device_name, dtype in itertools.product( + ("cpu", "gpu"), ("float64", "float32") + ): + yield array_namespace, device_name, dtype + + elif array_namespace == "array_api_strict": + # Always yield strings for consistent parametrization; _array_api_for_tests + # creates Device objects when needed. + yield array_namespace, "CPU_DEVICE", "float64" + yield array_namespace, "device1", "float32" + else: + yield array_namespace, None, None