Code review

Author: crusaderky

src/array_api_extra/_lib/_lazy.py

@@ -179,9 +179,9 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
     Raises
     ------
-    jax.errors.TracerArrayConversionError
-        When ``xp=jax.numpy``, `shape` is unknown (it contains None on one or more axes)
-        and this function was called inside ``jax.jit``.
+    ValueError
+        When ``xp=jax.numpy``, the output `shape` is unknown (it contains ``None`` on
+        one or more axes) and this function was called inside ``jax.jit``.
     RuntimeError
         When ``xp=sparse`` and auto-densification is disabled.
     Exception (backend-specific)
@@ -272,6 +272,10 @@ def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         # https://github.com/jax-ml/jax/issues/26102
         import jax
 
+        if any(None in shape for shape in shapes):
+            msg = "Output shape must be fully known when running inside jax.jit"
+            raise ValueError(msg)
+
         # Shield eager kwargs from being coerced into JAX arrays.
         # jax.pure_callback calls jax.jit under the hood, but without the chance of
         # passing static_argnames / static_argnums.

tests/test_lazy.py

@@ -191,26 +191,40 @@ def f(x: Array) -> Array:
     xp_assert_equal(y.compute(), x_cp + 1)  # type: ignore[attr-defined]  # pyright: ignore[reportUnknownArgumentType,reportAttributeAccessIssue]
 
 
-@pytest.mark.xfail_xp_backend(Backend.JAX, reason="unknown shape")
 def test_lazy_apply_none_shape_in_args(xp: ModuleType, library: Backend):
     x = xp.asarray([1, 1, 2, 2, 2])
 
-    xp2 = np if library is Backend.DASK else xp
-
-    # Single output
-    values = lazy_apply(xp2.unique_values, x, shape=(None,))
-    xp_assert_equal(values, xp.asarray([1, 2]))
-
-    # Multi output
+    # TODO mxp = meta_namespace(x, xp=xp)
+    mxp = np if library is Backend.DASK else xp
     int_type = xp.asarray(0).dtype
-    values, counts = lazy_apply(
-        xp2.unique_counts,
-        x,
-        shape=((None,), (None,)),
-        dtype=(x.dtype, int_type),
-    )
-    xp_assert_equal(values, xp.asarray([1, 2]))
-    xp_assert_equal(counts, xp.asarray([2, 3]))
+
+    if library is Backend.JAX:
+        # Single output
+        with pytest.raises(ValueError, match="Output shape must be fully known"):
+            _ = lazy_apply(mxp.unique_values, x, shape=(None,))
+
+        # Multi output
+        with pytest.raises(ValueError, match="Output shape must be fully known"):
+            _ = lazy_apply(
+                mxp.unique_counts,
+                x,
+                shape=((None,), (None,)),
+                dtype=(x.dtype, int_type),
+            )
+    else:
+        # Single output
+        values = lazy_apply(mxp.unique_values, x, shape=(None,))
+        xp_assert_equal(values, xp.asarray([1, 2]))
+
+        # Multi output
+        values, counts = lazy_apply(
+            mxp.unique_counts,
+            x,
+            shape=((None,), (None,)),
+            dtype=(x.dtype, int_type),
+        )
+        xp_assert_equal(values, xp.asarray([1, 2]))
+        xp_assert_equal(counts, xp.asarray([2, 3]))
 
 
 def check_lazy_apply_none_shape_broadcast(x: Array) -> Array:
@@ -349,10 +363,8 @@ def eager(
 def test_lazy_apply_kwargs(xp: ModuleType, library: Backend, as_numpy: bool):
     """When as_numpy=True, search and replace arrays in the (nested) keywords arguments
     with numpy arrays, and leave the rest untouched."""
-    expect_cls = (
-        np.ndarray if as_numpy or library is Backend.DASK else type(xp.asarray(0))
-    )
     x = xp.asarray(0)
+    expect_cls = np.ndarray if as_numpy or library is Backend.DASK else type(x)
     actual = check_lazy_apply_kwargs(x, expect_cls, as_numpy)  # pyright: ignore[reportUnknownArgumentType]
     xp_assert_equal(actual, x + 1)