dnm

Author: crusaderky

src/array_api_extra/testing.py

@@ -233,15 +233,20 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=PR01,RT0
 
         This allows `pytest.raises` and `pytest.warns` to work as expected.
         """
+        print(f"calling {func}")
         out = func(*args, **kwargs)
-
+        print(f"func returned {type(out)}")
+        out = jax.block_until_ready(out)
+        print(f"block_until_ready returned {type(out)}")
         # Note: `jax.block_until_ready` doesn't reraise exceptions.
         if isinstance(out, list | tuple):
             for o in out:
                 o.__array__()
         else:
+            print(f"collecting exceptions on {type(out)}")
             out.__array__()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
 
+        print(f"returning {type(out)}")
         return out  # pyright: ignore[reportReturnType]
 
     return wrapper