ENH: is_lazy_array()

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -819,6 +819,63 @@ def is_writeable_array(x) -> bool:
     return True
 
 
+def is_lazy_array(x) -> bool:
+    """Return True if x is potentially a future or it may be otherwise impossible or
+    expensive to eagerly read its contents, regardless of their size, e.g. by 
+    calling ``bool(x)`` or ``float(x)``.
+
+    Return True otherwise; e.g. ``bool(x)`` etc. is guaranteed to succeed and to be
+    cheap as long as the array is the right dtype.
+
+    Note
+    ----
+    This function errs on the side of caution for array types that may or may not be
+    lazy, e.g. JAX arrays, by always returning True for them.
+    """
+    if (
+        is_numpy_array(x) 
+        or is_cupy_array(x) 
+        or is_torch_array(x) 
+        or is_pydata_sparse_array(x)
+    ):
+        return False
+
+    # **JAX note:** while it is possible to determine if you're inside or outside
+    # jax.jit by testing the subclass of a jax.Array object, as well as testing bool()
+    # as we do below for unknown arrays, this is not recommended by JAX best practices.
+
+    # **Dask note:** Dask eagerly computes the graph on __bool__, __float__, and so on.
+    # This behaviour, while impossible to change without breaking backwards
+    # compatibility, is highly detrimental to performance as the whole graph will end
+    # up being computed multiple times.
+
+    if is_jax_array(x) or is_dask_array(x) or is_ndonnx_array(x):
+        return True
+
+    # Unknown Array API compatible object. Note that this test may have dire consequences
+    # in terms of performance, e.g. for a lazy object that eagerly computes the graph
+    # on __bool__ (dask is one such example, which however is special-cased above).
+
+    # Select a single point of the array
+    s = size(x)
+    if math.isnan(s):
+        return True
+    xp = array_namespace(x)
+    if s > 1:
+        x = xp.reshape(x, (-1,))[0]
+    # Cast to dtype=bool and deal with size 0 arrays
+    x = xp.any(x)
+
+    try:
+        bool(x)
+        return False
+    # The Array API standard dictactes that __bool__ should raise TypeError if the
+    # output cannot be defined.
+    # Here we're more lenient and also allow for e.g. NotImplementedError.
+    except Exception:
+        return True
+
+
 __all__ = [
     "array_namespace",
     "device",
@@ -840,6 +897,7 @@ def is_writeable_array(x) -> bool:
     "is_pydata_sparse_array",
     "is_pydata_sparse_namespace",
     "is_writeable_array",
+    "is_lazy_array",
     "size",
     "to_device",
 ]

docs/helper-functions.rst

@@ -52,6 +52,7 @@ yet.
 .. autofunction:: is_pydata_sparse_array
 .. autofunction:: is_ndonnx_array
 .. autofunction:: is_writeable_array
+.. autofunction:: is_lazy_array
 .. autofunction:: is_numpy_namespace
 .. autofunction:: is_cupy_namespace
 .. autofunction:: is_torch_namespace

tests/test_common.py

@@ -5,8 +5,9 @@
     is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace,
 )
 
-from array_api_compat import device, is_array_api_obj, is_writeable_array, to_device
-
+from array_api_compat import (
+    device, is_array_api_obj, is_lazy_array, is_writeable_array, to_device
+)
 from ._helpers import import_, wrapped_libraries, all_libraries
 
 import pytest
@@ -92,6 +93,70 @@ def test_is_writeable_array_numpy():
     assert not is_writeable_array(x)
 
 
+@pytest.mark.parametrize("library", all_libraries)
+def test_is_lazy_array(library):
+    lib = import_(library)
+    x = lib.asarray([1, 2, 3])
+    assert isinstance(is_lazy_array(x), bool)
+
+
+@pytest.mark.parametrize("array", [
+    [], [1, 2], 1, 0, float("nan"), [[1, 2], [3, 4]]
+])
+def test_is_lazy_array_unknown(array, monkeypatch):
+    """Test is_lazy_array() on an unknown Array API compliant object"""
+    xp = import_("jax.numpy")
+    import array_api_compat.common._helpers
+    import jax
+
+    x = xp.asarray(array)
+    # Prevent is_lazy_array() from special-casing JAX
+    monkeypatch.setattr(
+        array_api_compat.common._helpers, 
+        "is_jax_array", 
+        lambda x: False,
+    )
+
+    assert not is_lazy_array(x)  # Eager JAX
+    assert jax.jit(is_lazy_array)(x)  # Jitted (lazy) JAX
+
+
+def test_is_lazy_array_unknown_dask(monkeypatch):
+    """Test is_lazy_array() on an unknown Array API compliant object which
+    - may or may not raise an arbitrary exception on bool()
+    - may or may not have NaN in its shape
+    """
+    da = import_("dask.array", wrapper=True)
+    import array_api_compat.common._helpers
+
+    x = da.arange(10)
+    y = x[x > 5]
+    assert np.isnan(y.size)
+
+    def do_not_run(_):
+        raise AssertionError("do_not_run")
+
+    z = x.map_blocks(do_not_run, dtype=x.dtype)
+    with pytest.raises(AssertionError, match="do_not_run"):
+        z.compute()
+
+    # Prevent is_lazy_array() from special-casing Dask
+    monkeypatch.setattr(
+        array_api_compat.common._helpers, 
+        "is_dask_array", 
+        lambda x: False,
+    )
+    monkeypatch.setattr(
+        array_api_compat.common._helpers, 
+        "array_namespace", 
+        lambda x: da,
+    )
+
+    assert not is_lazy_array(x)  # Eagerly computes on bool()
+    assert is_lazy_array(y)  # NaN size
+    assert is_lazy_array(z)  # bool() raises AssertionError
+
+
 @pytest.mark.parametrize("library", all_libraries)
 def test_device(library):
     xp = import_(library, wrapper=True)
@@ -149,6 +214,7 @@ def test_asarray_cross_library(source_library, target_library, request):
 
     assert is_tgt_type(b), f"Expected {b} to be a {tgt_lib.ndarray}, but was {type(b)}"
 
+
 @pytest.mark.parametrize("library", wrapped_libraries)
 def test_asarray_copy(library):
     # Note, we have this test here because the test suite currently doesn't