Recursion guard

Author: crusaderky

src/array_api_extra/_lib/_lazy.py

@@ -324,28 +324,48 @@ def _contains_jax_arrays(x: object) -> bool:  # numpydoc ignore=PR01,RT01
     """
     Test if x is a JAX array or a nested collection with any JAX arrays in it.
     """
-    if is_jax_array(x):
-        return True
-    if isinstance(x, list | tuple):
-        return any(_contains_jax_arrays(i) for i in x)  # pyright: ignore[reportUnknownArgumentType]
-    if isinstance(x, dict):
-        return any(_contains_jax_arrays(i) for i in x.values())  # pyright: ignore[reportUnknownArgumentType]
-    return False
+    seen = set()
+
+    def recursion(x: object) -> bool:  # numpydoc ignore=GL08
+        if id(x) in seen:
+            return False
+        seen.add(id(x))
+
+        if is_jax_array(x):
+            return True
+        if isinstance(x, list | tuple):
+            return any(recursion(i) for i in x)  # pyright: ignore[reportUnknownArgumentType]
+        if isinstance(x, dict):
+            return any(recursion(i) for i in x.values())  # pyright: ignore[reportUnknownArgumentType]
+        return False
+
+    return recursion(x)
 
 
 def _as_numpy(x: object) -> Any:  # type: ignore[no-any-explicit] # numpydoc ignore=PR01,RT01
     """Recursively convert Array API objects in x to NumPy."""
     import numpy as np  # pylint: disable=import-outside-toplevel
 
-    if is_array_api_obj(x):
-        return np.asarray(x)
-    if isinstance(x, list) or type(x) is tuple:  # pylint: disable=unidiomatic-typecheck
-        return type(x)(_as_numpy(i) for i in x)  # pyright: ignore[reportUnknownArgumentType]
-    if isinstance(x, tuple):  # namedtuple
-        return type(x)(*(_as_numpy(i) for i in x))  # pyright: ignore[reportUnknownArgumentType]
-    if isinstance(x, dict):
-        return {k: _as_numpy(v) for k, v in x.items()}  # pyright: ignore[reportUnknownArgumentType]
-    return x
+    seen = set()
+
+    def recursion(x: Any) -> Any:  # type: ignore[no-any-explicit]  # numpydoc ignore=GL08
+        if is_array_api_obj(x):
+            return np.asarray(x)
+        if not isinstance(x, list | tuple | dict):
+            return x
+
+        if id(x) in seen:  # pyright: ignore[reportUnknownArgumentType]
+            return x  # Recursive collections can't contain arrays
+        seen.add(id(x))  # pyright: ignore[reportUnknownArgumentType]
+
+        if isinstance(x, list) or type(x) is tuple:  # pylint: disable=unidiomatic-typecheck  # pyright: ignore[reportUnknownArgumentType]
+            return type(x)(recursion(i) for i in x)  # pyright: ignore[reportUnknownArgumentType]
+        if isinstance(x, tuple):  # namedtuple
+            return type(x)(*(recursion(i) for i in x))  # pyright: ignore[reportUnknownArgumentType]
+        # dict
+        return {k: recursion(v) for k, v in x.items()}
+
+    return recursion(x)
 
 
 def _lazy_apply_wrapper(  # type: ignore[no-any-explicit]  # numpydoc ignore=PR01,RT01

tests/test_lazy.py

@@ -9,7 +9,7 @@
 from array_api_extra._lib import Backend
 from array_api_extra._lib._testing import xp_assert_equal
 from array_api_extra._lib._utils import _compat
-from array_api_extra._lib._utils._compat import array_namespace
+from array_api_extra._lib._utils._compat import array_namespace, is_dask_array
 from array_api_extra._lib._utils._typing import Array, Device
 from array_api_extra.testing import lazy_xp_function
 
@@ -288,12 +288,18 @@ class NT(NamedTuple):
 
 
 def check_lazy_apply_kwargs(x: Array, expect_cls: type, as_numpy: bool) -> Array:
+    is_dask = is_dask_array(x)
+    recursive: list[object] = []
+    if not is_dask:  # dask.delayed crashes on recursion
+        recursive.append(recursive)
+
     def eager(
         x: Array,
         z: dict[str, list[Array] | tuple[Array, ...] | NT],
         msg: str,
         msgs: list[str],
         scalar: int,
+        recursive: list[list[object]],
     ) -> Array:
         assert isinstance(x, expect_cls)
         # JAX will crash if x isn't material
@@ -310,6 +316,9 @@ def eager(
         assert isinstance(msgs[0], str)
         assert scalar == 1  # must be hidden from JAX
         assert isinstance(scalar, int)
+        assert isinstance(recursive, list)
+        if not is_dask:
+            assert recursive[0][0] is recursive[0]
         return x + 1  # type: ignore[operator]
 
     # Use explicit namespace to bypass monkey-patching by lazy_xp_function
@@ -323,6 +332,7 @@ def eager(
         msgs=["Hello World"],
         # This will be automatically cast to jax.Array if we don't wrap it
         scalar=1,
+        recursive=recursive,
         shape=x.shape,
         dtype=x.dtype,
         as_numpy=as_numpy,