Test none shapes

Author: crusaderky

Diff for: src/array_api_extra/_lib/_testing.py

@@ -4,11 +4,13 @@
 Note that this is private API; don't expect it to be stable.
 """
 
+import math
 from types import ModuleType
 
 from ._utils._compat import (
     array_namespace,
     is_cupy_namespace,
+    is_dask_namespace,
     is_pydata_sparse_namespace,
     is_torch_namespace,
 )
@@ -40,8 +42,16 @@ def _check_ns_shape_dtype(
     msg = f"namespaces do not match: {actual_xp} != f{desired_xp}"
     assert actual_xp == desired_xp, msg
 
-    msg = f"shapes do not match: {actual.shape} != f{desired.shape}"
-    assert actual.shape == desired.shape, msg
+    actual_shape = actual.shape
+    desired_shape = desired.shape
+    if is_dask_namespace(desired_xp):
+        if any(math.isnan(i) for i in actual_shape):
+            actual_shape = actual.compute().shape
+        if any(math.isnan(i) for i in desired_shape):
+            desired_shape = desired.compute().shape
+
+    msg = f"shapes do not match: {actual_shape} != f{desired_shape}"
+    assert actual_shape == desired_shape, msg
 
     msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
     assert actual.dtype == desired.dtype, msg

Diff for: tests/test_funcs.py

@@ -366,6 +366,20 @@ def test_bool_dtype(self, xp: ModuleType):
             isclose(xp.asarray(True), b, atol=1), xp.asarray([True, True, True])
         )
 
+    def test_none_shape(self, xp: ModuleType):
+        a = xp.asarray([1, 5, 0])
+        b = xp.asarray([1, 4, 2])
+        b = b[a < 5]
+        a = a[a < 5]
+        xp_assert_equal(isclose(a, b), xp.asarray([True, False]))
+
+    def test_none_shape_bool(self, xp: ModuleType):
+        a = xp.asarray([True, True, False])
+        b = xp.asarray([True, False, True])
+        b = b[a]
+        a = a[a]
+        xp_assert_equal(isclose(a, b), xp.asarray([True, False]))
+
     def test_xp(self, xp: ModuleType):
         a = xp.asarray([0.0, 0.0])
         b = xp.asarray([1e-9, 1e-4])

Diff for: tests/test_testing.py

@@ -10,8 +10,7 @@
 # mypy: disable-error-code=no-any-decorated
 # pyright: reportUnknownParameterType=false,reportMissingParameterType=false
 
-
-@pytest.mark.parametrize(
+param_assert_equal_close = pytest.mark.parametrize(
     "func",
     [
         xp_assert_equal,
@@ -21,6 +20,9 @@
         ),
     ],
 )
+
+
+@param_assert_equal_close
 def test_assert_close_equal_basic(xp: ModuleType, func: Callable[..., None]):  # type: ignore[no-any-explicit]
     func(xp.asarray(0), xp.asarray(0))
     func(xp.asarray([1, 2]), xp.asarray([1, 2]))
@@ -68,3 +70,21 @@ def test_assert_close_tolerance(xp: ModuleType):
     xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), atol=3)
     with pytest.raises(AssertionError):
         xp_assert_close(xp.asarray([100.0]), xp.asarray([102.0]), atol=1)
+
+
+@param_assert_equal_close
+@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="no bool indexing by sparse arrays")
+def test_assert_close_equal_none_shape(xp: ModuleType, func: Callable[..., None]):  # type: ignore[no-any-explicit]
+    """On dask and other lazy backends, test that a shape with NaN's or None's
+    can be compared to a real shape.
+    """
+    a = xp.asarray([1, 2])
+    a = a[a > 1]
+
+    func(a, xp.asarray([2]))
+    with pytest.raises(AssertionError):
+        func(a, xp.asarray([2, 3]))
+    with pytest.raises(AssertionError):
+        func(a, xp.asarray(2))
+    with pytest.raises(AssertionError):
+        func(a, xp.asarray([3]))