diff --git a/src/array_api_extra/testing.py b/src/array_api_extra/testing.py index 6b46e09f..42ad6a43 100644 --- a/src/array_api_extra/testing.py +++ b/src/array_api_extra/testing.py @@ -28,6 +28,9 @@ is_torch_namespace, to_device, ) +from ._lib._utils._compat import ( + device as get_device, +) from ._lib._utils._helpers import jax_autojit, pickle_flatten, pickle_unflatten from ._lib._utils._typing import Array, Device @@ -564,12 +567,13 @@ def _require_numpy() -> ModuleType: # numpydoc ignore=RT01 return np -def _check_ns_shape_dtype( +def _check_ns_shape_dtype_device( actual: Array, desired: Array, check_dtype: bool, check_shape: bool, check_scalar: bool, + check_device: bool, xp: ModuleType | None = None, ) -> tuple[Array, Array, ModuleType, ModuleType]: # numpydoc ignore=RT03 """ @@ -588,6 +592,8 @@ def _check_ns_shape_dtype( check_scalar : bool, default: False NumPy only: whether to check agreement between actual and desired types - 0d array vs scalar. + check_device : bool, default: True + Whether to check agreement between actual and desired devices. xp : array_namespace, optional A standard-compatible namespace which `actual` and `desired` must match. @@ -655,6 +661,12 @@ def _check_ns_shape_dtype( if check_dtype: msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}" assert actual.dtype == desired.dtype, msg + if check_device: + msg = ( + f"Devices do not match.\nActual: {get_device(actual)}\n" + f"Desired: {get_device(desired)}" + ) + assert get_device(actual) == get_device(desired), msg desired = desired_xp.broadcast_to(desired, actual_shape) return actual, desired, desired_xp, np @@ -714,6 +726,7 @@ def assert_close( check_dtype: bool = True, check_shape: bool = True, check_scalar: bool = False, + check_device: bool = True, xp: ModuleType | None = None, ) -> None: """ @@ -746,6 +759,8 @@ def assert_close( 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`). + check_device : bool, default: True + Whether to check agreement between actual and desired devices. xp : array_namespace, optional A standard-compatible namespace which `actual` and `desired` must match. @@ -777,8 +792,8 @@ def assert_close( Array arguments to `atol` and `rtol` must be valid input to :class:`float`. """ __tracebackhide__ = True - actual, desired, xp, np = _check_ns_shape_dtype( - actual, desired, check_dtype, check_shape, check_scalar, xp + actual, desired, xp, np = _check_ns_shape_dtype_device( + actual, desired, check_dtype, check_shape, check_scalar, check_device, xp ) if not _is_materializable(actual): return @@ -818,6 +833,7 @@ def assert_equal( check_dtype: bool = True, check_shape: bool = True, check_scalar: bool = False, + check_device: bool = True, xp: ModuleType | None = None, ) -> None: """ @@ -844,6 +860,8 @@ def assert_equal( 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`). + check_device : bool, default: True + Whether to check agreement between actual and desired devices. xp : array_namespace, optional A standard-compatible namespace which `actual` and `desired` must match. @@ -861,8 +879,8 @@ def assert_equal( numpy.testing.assert_array_equal : Similar function for NumPy arrays. """ __tracebackhide__ = True - actual, desired, xp, np = _check_ns_shape_dtype( - actual, desired, check_dtype, check_shape, check_scalar, xp + actual, desired, xp, np = _check_ns_shape_dtype_device( + actual, desired, check_dtype, check_shape, check_scalar, check_device, xp ) if not _is_materializable(actual): return @@ -882,6 +900,7 @@ def assert_less( check_dtype: bool = True, check_shape: bool = True, check_scalar: bool = False, + check_device: bool = True, xp: ModuleType | None = None, ) -> None: """ @@ -906,6 +925,8 @@ def assert_less( 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`). + check_device : bool, default: True + Whether to check agreement between actual and desired devices. xp : array_namespace, optional A standard-compatible namespace which `x` and `y` must match. @@ -923,8 +944,8 @@ def assert_less( numpy.testing.assert_array_less : Similar function for NumPy arrays. """ __tracebackhide__ = True - x, y, xp, np = _check_ns_shape_dtype( - x, y, check_dtype, check_shape, check_scalar, xp + x, y, xp, np = _check_ns_shape_dtype_device( + x, y, check_dtype, check_shape, check_scalar, check_device, xp ) if not _is_materializable(x): return @@ -941,6 +962,7 @@ def assert_close_nulp( check_dtype: bool = True, check_shape: bool = True, check_scalar: bool = False, + check_device: bool = True, xp: ModuleType | None = None, ) -> None: """ @@ -966,6 +988,8 @@ def assert_close_nulp( 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`). + check_device : bool, default: True + Whether to check agreement between actual and desired devices. xp : array_namespace, optional A standard-compatible namespace which `actual` and `desired` must match. @@ -996,8 +1020,8 @@ def assert_close_nulp( 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 + actual, desired, xp, np = _check_ns_shape_dtype_device( + actual, desired, check_dtype, check_shape, check_scalar, check_device, xp ) if not _is_materializable(actual): return diff --git a/tests/test_funcs.py b/tests/test_funcs.py index c4a0fee5..5114c61b 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -211,15 +211,6 @@ def test_xp(self, xp: ModuleType): expect = xp.where(cond, self.f1(x), self.f2(x)) assert_equal(actual, expect) - def test_device(self, xp: ModuleType, device: Device): - x = xp.asarray([1, 2, 3, 4], device=device) - y = apply_where(x % 2 == 0, x, self.f1, self.f2) - assert get_device(y) == device - y = apply_where(x % 2 == 0, x, self.f1, fill_value=0) - assert get_device(y) == device - y = apply_where(x % 2 == 0, x, self.f1, fill_value=x) - assert get_device(y) == device - @pytest.mark.filterwarnings("ignore::RuntimeWarning") # overflows, etc. @hypothesis.settings( # The xp and library fixtures are not regenerated between hypothesis iterations @@ -1050,19 +1041,6 @@ def test_device(self, xp: ModuleType, device: Device, equal_nan: bool): res = isclose(a, b, equal_nan=equal_nan) assert get_device(res) == device - def test_array_on_device_with_scalar(self, xp: ModuleType, device: Device): - a = xp.asarray([0.01, 0.5, 0.8, 0.9, 1.00001], device=device, dtype=xp.float64) - b = 1 - res = isclose(a, b) - assert get_device(res) == device - assert_equal(res, xp.asarray([False, False, False, False, True])) - - a = 0.1 - b = xp.asarray([0.01, 0.5, 0.8, 0.9, 0.100001], device=device, dtype=xp.float64) - res = isclose(a, b) - assert get_device(res) == device - assert_equal(res, xp.asarray([False, False, False, False, True])) - class TestKron: def test_basic(self, xp: ModuleType): @@ -1668,17 +1646,14 @@ def test_device(self, xp: ModuleType, device: Device, library: Backend): b = xp.asarray([1, 2, 3], device=device) assert get_device(isin(a, b)) == device - def test_assume_unique_and_invert( - self, xp: ModuleType, device: Device, library: Backend - ): + def test_assume_unique_and_invert(self, xp: ModuleType, library: Backend): if library.like(Backend.NUMPY) and NUMPY_VERSION < (1, 24): pytest.xfail("NumPy <1.24 has no kind kwarg in isin") - a = xp.asarray([0, 3, 6, 10], device=device) - b = xp.asarray([1, 2, 3, 10], device=device) + a = xp.asarray([0, 3, 6, 10]) + b = xp.asarray([1, 2, 3, 10]) expected = xp.asarray([True, False, True, False]) res = isin(a, b, assume_unique=True, invert=True) - assert get_device(res) == device assert_equal(res, expected) def test_kind(self, xp: ModuleType, library: Backend): diff --git a/tests/test_testing.py b/tests/test_testing.py index 84ee5ba3..77322f03 100644 --- a/tests/test_testing.py +++ b/tests/test_testing.py @@ -13,6 +13,9 @@ is_dask_namespace, is_jax_namespace, ) +from array_api_extra._lib._utils._compat import ( + device as get_device, +) from array_api_extra._lib._utils._typing import Array, Device from array_api_extra.testing import ( _as_numpy_array, @@ -206,13 +209,25 @@ def test_device(self, xp: ModuleType, device: Device, func: Callable[..., None]) b = xp.asarray([2], device=device) c = xp.asarray([2, 2], device=device) - func(a, b) + func(a, b, check_device=True) with pytest.raises(AssertionError, match="shapes do not match"): func(a, c) # This is normally performed by np.testing.assert_array_equal etc. # but in case of torch device='meta' we have to do it manually with pytest.raises(AssertionError, match="sizes do not match"): - func(a, c, check_shape=False) + func(a, c, check_shape=False, check_device=False) + + # Checks for non-default device + default_device = get_device(xp.empty(0)) + if device != default_device: + d = xp.asarray([2], device=default_device) + with pytest.raises(AssertionError, match="Devices do not match"): + func(a, d) + func(a, d, check_device=False) + if getattr(device, "type", None) != "meta": + e = xp.asarray([0], device=default_device) + with pytest.raises(AssertionError, match=self.np_err_msg(func)): + func(a, e, check_device=False) def test_assert_close_nulp(self, xp: ModuleType): a = xp.asarray([1.0, 1e-10], dtype=xp.float64)