test eager_shape

Author: crusaderky

tests/test_helpers.py

@@ -7,7 +7,7 @@
 from array_api_extra._lib import Backend
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
-from array_api_extra._lib._utils._helpers import asarrays, in1d, ndindex
+from array_api_extra._lib._utils._helpers import asarrays, eager_shape, in1d, ndindex
 from array_api_extra._lib._utils._typing import Array, Device, DType
 from array_api_extra.testing import lazy_xp_function
 
@@ -156,3 +156,20 @@ def test_numpy_generics(self, dtype: DType):
 )
 def test_ndindex(shape: tuple[int, ...]):
     assert tuple(ndindex(*shape)) == tuple(np.ndindex(*shape))
+
+
+@pytest.mark.skip_xp_backend(Backend.SPARSE, reason="index by sparse array")
+def test_eager_shape(xp: ModuleType, library: Backend):
+    a = xp.asarray([1, 2, 3])
+    # Lazy arrays, like Dask, have an eager shape until you slice them with
+    # a lazy boolean mask
+    assert eager_shape(a) == a.shape == (3,)
+
+    b = a[a > 2]
+    if library is Backend.DASK:
+        with pytest.raises(TypeError, match="Unsupported lazy shape"):
+            _ = eager_shape(b)
+    # FIXME can't test use case for None in the shape until we add support for
+    # other lazy backends
+    else:
+        assert eager_shape(b) == b.shape == (1,)