Fix xp_array_less

Author: crusaderky

src/array_api_extra/_lib/_testing.py

@@ -22,6 +22,7 @@
     is_jax_namespace,
     is_numpy_namespace,
     is_pydata_sparse_namespace,
+    is_torch_array,
     is_torch_namespace,
     to_device,
 )
@@ -62,18 +63,28 @@ def _check_ns_shape_dtype(
     msg = f"namespaces do not match: {actual_xp} != f{desired_xp}"
     assert actual_xp == desired_xp, msg
 
-    if check_shape:
-        actual_shape = actual.shape
-        desired_shape = desired.shape
-        if is_dask_namespace(desired_xp):
-            # Dask uses nan instead of None for unknown shapes
-            if any(math.isnan(i) for i in cast(tuple[float, ...], actual_shape)):
-                actual_shape = actual.compute().shape  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
-            if any(math.isnan(i) for i in cast(tuple[float, ...], desired_shape)):
-                desired_shape = desired.compute().shape  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+    # Dask uses nan instead of None for unknown shapes
+    actual_shape = cast(tuple[float, ...], actual.shape)
+    desired_shape = cast(tuple[float, ...], desired.shape)
+    assert None not in actual_shape  # Requires explicit support
+    assert None not in desired_shape
+    if is_dask_namespace(desired_xp):
+        if any(math.isnan(i) for i in actual_shape):
+            actual_shape = actual.compute().shape  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+        if any(math.isnan(i) for i in desired_shape):
+            desired_shape = desired.compute().shape  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
 
+    if check_shape:
         msg = f"shapes do not match: {actual_shape} != f{desired_shape}"
         assert actual_shape == desired_shape, msg
+    else:
+        # Ignore shape, but check flattened size. This is normally done by
+        # np.testing.assert_array_equal etc even when strict=False, but not for
+        # non-materializable arrays.
+        actual_size = math.prod(actual_shape)  # pyright: ignore[reportUnknownArgumentType]
+        desired_size = math.prod(desired_shape)  # pyright: ignore[reportUnknownArgumentType]
+        msg = f"sizes do not match: {actual_size} != f{desired_size}"
+        assert actual_size == desired_size, msg
 
     if check_dtype:
         msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
@@ -90,6 +101,17 @@ def _check_ns_shape_dtype(
     return desired_xp
 
 
+def _is_materializable(x: Array) -> bool:
+    """
+    Check if the array is materializable, e.g. `as_numpy_array` can be called on it
+    and one can assume that `__dlpack__` will succeed (if implemented, and given a
+    compatible device).
+    """
+    # Important: here we assume that we're not tracing -
+    # e.g. we're not inside `jax.jit`` nor `cupy.cuda.Stream.begin_capture`.
+    return not is_torch_array(x) or x.device.type != "meta"  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+
+
 def as_numpy_array(array: Array, *, xp: ModuleType) -> np.typing.NDArray[Any]:  # type: ignore[explicit-any]
     """
     Convert array to NumPy, bypassing GPU-CPU transfer guards and densification guards.
@@ -100,11 +122,7 @@ def as_numpy_array(array: Array, *, xp: ModuleType) -> np.typing.NDArray[Any]:
         return array.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
 
     if is_torch_namespace(xp):
-        if array.device.type == "meta":  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
-            # Can't materialize; generate dummy data instead
-            array = xp.zeros_like(array, device="cpu")
-        else:
-            array = to_device(array, "cpu")
+        array = to_device(array, "cpu")
     if is_array_api_strict_namespace(xp):
         cpu: Device = xp.Device("CPU_DEVICE")
         array = to_device(array, cpu)
@@ -150,6 +168,8 @@ def xp_assert_equal(
     numpy.testing.assert_array_equal : Similar function for NumPy arrays.
     """
     xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
+    if not _is_materializable(actual):
+        return
     actual_np = as_numpy_array(actual, xp=xp)
     desired_np = as_numpy_array(desired, xp=xp)
     np.testing.assert_array_equal(actual_np, desired_np, err_msg=err_msg)
@@ -185,6 +205,8 @@ def xp_assert_less(
     numpy.testing.assert_array_equal : Similar function for NumPy arrays.
     """
     xp = _check_ns_shape_dtype(x, y, check_dtype, check_shape, check_scalar)
+    if not _is_materializable(x):
+        return
     x_np = as_numpy_array(x, xp=xp)
     y_np = as_numpy_array(y, xp=xp)
     np.testing.assert_array_less(x_np, y_np, err_msg=err_msg)
@@ -233,6 +255,8 @@ def xp_assert_close(
     The default `atol` and `rtol` differ from `xp.all(xpx.isclose(a, b))`.
     """
     xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
+    if not _is_materializable(actual):
+        return
 
     if rtol is None:
         if xp.isdtype(actual.dtype, ("real floating", "complex floating")):

tests/test_testing.py

@@ -30,16 +30,11 @@ def test_basic(self, xp: ModuleType):
         y = as_numpy_array(x, xp=xp)
         xp_assert_equal(y, np.asarray([1, 2, 3]))  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
-    def test_device(self, xp: ModuleType, library: Backend, device: Device):
+    @pytest.mark.skip_xp_backend(Backend.TORCH, reason="materialize 'meta' device")
+    def test_device(self, xp: ModuleType, device: Device):
         x = xp.asarray([1, 2, 3], device=device)
-        actual = as_numpy_array(x, xp=xp)
-        if library is Backend.TORCH:
-            assert device.type == "meta"  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
-            expect = np.asarray([0, 0, 0])
-        else:
-            expect = np.asarray([1, 2, 3])
-
-        xp_assert_equal(actual, expect)  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
+        y = as_numpy_array(x, xp=xp)
+        xp_assert_equal(y, np.asarray([1, 2, 3]))  # type: ignore[arg-type]  # pyright: ignore[reportArgumentType]
 
 
 class TestAssertEqualCloseLess:
@@ -92,7 +87,7 @@ def test_check_shape(self, xp: ModuleType, func: Callable[..., None]):
         func(a, b, check_shape=False)
         with pytest.raises(AssertionError, match="Mismatched elements"):
             func(a, c, check_shape=False)
-        with pytest.raises(AssertionError, match=r"shapes \(1,\), \(2,\) mismatch"):
+        with pytest.raises(AssertionError, match="sizes do not match"):
             func(a, d, check_shape=False)
 
     @pytest.mark.parametrize("func", [xp_assert_equal, pr_assert_close, xp_assert_less])
@@ -181,6 +176,20 @@ def test_none_shape(self, xp: ModuleType, func: Callable[..., None]):
         with pytest.raises(AssertionError, match="Mismatched elements"):
             func(xp.asarray([4]), a)
 
+    @pytest.mark.parametrize("func", [xp_assert_equal, pr_assert_close, xp_assert_less])
+    def test_device(self, xp: ModuleType, device: Device, func: Callable[..., None]):
+        a = xp.asarray([1] if func is xp_assert_less else [2], device=device)
+        b = xp.asarray([2], device=device)
+        c = xp.asarray([2, 2], device=device)
+
+        func(a, b)
+        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)
+
 
 def good_lazy(x: Array) -> Array:
     """A function that behaves well in Dask and jax.jit"""