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..19fd5b37 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -1055,13 +1055,13 @@ def test_array_on_device_with_scalar(self, xp: ModuleType, device: Device): b = 1 res = isclose(a, b) assert get_device(res) == device - assert_equal(res, xp.asarray([False, False, False, False, True])) + assert_equal(res, xp.asarray([False, False, False, False, True], device=device)) 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])) + assert_equal(res, xp.asarray([False, False, False, False, True], device=device)) class TestKron: @@ -1676,7 +1676,7 @@ def test_assume_unique_and_invert( a = xp.asarray([0, 3, 6, 10], device=device) b = xp.asarray([1, 2, 3, 10], device=device) - expected = xp.asarray([True, False, True, False]) + expected = xp.asarray([True, False, True, False], device=device) res = isin(a, b, assume_unique=True, invert=True) assert get_device(res) == device assert_equal(res, expected)