Support unknown shapes

Author: crusaderky

docs/api-reference.md

@@ -7,6 +7,7 @@
     :toctree: generated
 
     apply_numpy_func
+    UnknownShapeError
     at
     atleast_nd
     cov

src/array_api_extra/__init__.py

@@ -1,7 +1,7 @@
 """Extra array functions built on top of the array API standard."""
 
 from ._delegation import pad
-from ._lib._apply import apply_numpy_func
+from ._lib._apply import UnknownShapeError, apply_numpy_func
 from ._lib._funcs import (
     at,
     atleast_nd,
@@ -18,6 +18,7 @@
 
 # pylint: disable=duplicate-code
 __all__ = [
+    "UnknownShapeError",
     "__version__",
     "apply_numpy_func",
     "at",

src/array_api_extra/_lib/_apply.py

@@ -3,6 +3,7 @@
 # https://github.com/scikit-learn/scikit-learn/pull/27910#issuecomment-2568023972
 from __future__ import annotations
 
+import math
 from collections.abc import Callable, Sequence
 from functools import wraps
 from types import ModuleType
@@ -30,11 +31,19 @@ class P:  # pylint: disable=missing-class-docstring
         kwargs: dict
 
 
+class UnknownShapeError(ValueError):
+    """
+    `shape` contains one or more None elements.
+
+    This is unsupported when running inside `jax.jit`.
+    """
+
+
 @overload
 def apply_numpy_func(  # type: ignore[valid-type]
     func: Callable[P, NumPyObject],
     *args: Array,
-    shape: tuple[int, ...] | None = None,
+    shape: tuple[int | None, ...] | None = None,
     dtype: DType | None = None,
     xp: ModuleType | None = None,
     **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
@@ -45,7 +54,7 @@ def apply_numpy_func(  # type: ignore[valid-type]
 def apply_numpy_func(  # type: ignore[valid-type]
     func: Callable[P, Sequence[NumPyObject]],
     *args: Array,
-    shape: Sequence[tuple[int, ...]],
+    shape: Sequence[tuple[int | None, ...]],
     dtype: Sequence[DType] | None = None,
     xp: ModuleType | None = None,
     **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
@@ -55,7 +64,7 @@ def apply_numpy_func(  # type: ignore[valid-type]
 def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     func: Callable[P, NumPyObject | Sequence[NumPyObject]],
     *args: Array,
-    shape: tuple[int, ...] | Sequence[tuple[int, ...]] | None = None,
+    shape: tuple[int | None, ...] | Sequence[tuple[int | None, ...]] | None = None,
     dtype: DType | Sequence[DType] | None = None,
     xp: ModuleType | None = None,
     **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
@@ -76,7 +85,7 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         One or more Array API compliant arrays. You need to be able to apply
         :func:`numpy.asarray` to them to convert them to numpy; read notes below about
         specific backends.
-    shape : tuple[int, ...] | Sequence[tuple[int, ...]], optional
+    shape : tuple[int | None, ...] | Sequence[tuple[int, ...]], optional
         Output shape or sequence of output shapes, one for each output of `func`.
         Default: assume single output and broadcast shapes of the input arrays.
     dtype : DType | Sequence[DType], optional
@@ -102,6 +111,8 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
     JAX
         This allows applying eager functions to jitted JAX arrays, which are lazy.
         The function won't be applied until the JAX array is materialized.
+        When running inside `jax.jit`, `shape` must be fully known, i.e. it cannot
+        contain any `None` elements.
 
         The :doc:`jax:transfer_guard` may prevent arrays on a GPU device from being
         transferred back to CPU. This is treated as an implicit transfer.
@@ -135,6 +146,18 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         :func:`dask.array.blockwise`, or a native Dask wrapper instead of
         `apply_numpy_func`.
 
+    Raises
+    ------
+    UnknownShapeError
+        When `shape` is unknown (one or more sizes are None) and this function was
+        called inside `jax.jit`.
+
+    Exception (varies)
+
+        - When the backend disallows implicit device to host transfers and the input
+          arrays are on a device, e.g. on GPU;
+        - When the backend is sparse and auto-densification is disabled.
+
     See Also
     --------
     jax.transfer_guard
@@ -147,13 +170,16 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         xp = array_namespace(*args)
 
     # Normalize and validate shape and dtype
+    shapes: list[tuple[int | None, ...]]
+    dtypes: list[DType]
     multi_output = False
+
     if shape is None:
         shapes = [xp.broadcast_shapes(*(arg.shape for arg in args))]
-    elif isinstance(shape, tuple) and all(isinstance(s, int) for s in shape):
-        shapes = [shape]
+    elif isinstance(shape, tuple) and all(isinstance(s, int | None) for s in shape):
+        shapes = [shape]  # pyright: ignore[reportAssignmentType]
     else:
-        shapes = list(shape)
+        shapes = list(shape)  # type: ignore[arg-type]  # pyright: ignore[reportAssignmentType]
         multi_output = True
 
     if dtype is None:
@@ -186,13 +212,19 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         meta_xp = array_namespace(*metas)
 
         wrapped = dask.delayed(_npfunc_wrapper(func, multi_output, meta_xp), pure=True)
-        # This finalizes each arg, which is the same as arg.rechunk(-1)
+        # This finalizes each arg, which is the same as arg.rechunk(-1).
         # Please read docstring above for why we're not using
         # dask.array.map_blocks or dask.array.blockwise!
         delayed_out = wrapped(*args, **kwargs)
 
         out = tuple(
-            xp.from_delayed(delayed_out[i], shape=shape, dtype=dtype, meta=metas[0])
+            xp.from_delayed(
+                delayed_out[i],
+                # Dask's unknown shapes diverge from the Array API specification
+                shape=tuple(math.nan if s is None else s for s in shape),
+                dtype=dtype,
+                meta=metas[0],
+            )
             for i, (shape, dtype) in enumerate(zip(shapes, dtypes, strict=True))
         )
 
@@ -205,18 +237,33 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         import jax  # type: ignore[import-not-found]  # pylint: disable=import-outside-toplevel,import-error  # pyright: ignore[reportMissingImports]
 
         wrapped = _npfunc_wrapper(func, multi_output, xp)
-        out = cast(
-            tuple[Array, ...],
-            jax.pure_callback(
-                wrapped,
-                tuple(
-                    jax.ShapeDtypeStruct(s, dt)  # pyright: ignore[reportUnknownArgumentType]
-                    for s, dt in zip(shapes, dtypes, strict=True)
+
+        if any(s is None for shape in shapes for s in shape):
+            # Unknown output shape. Won't work with jax.jit, but it
+            # can work with eager jax.
+            try:
+                out = wrapped(*args, **kwargs)
+            except jax.errors.TracerArrayConversionError:
+                msg = (
+                    "jax.jit can't delay application of numpy functions when the shape "
+                    "of the returned array(s) is unknown. "
+                    f"shape={shapes if multi_output else shapes[0]}"
+                )
+                raise UnknownShapeError(msg) from None
+
+        else:
+            out = cast(
+                tuple[Array, ...],
+                jax.pure_callback(
+                    wrapped,
+                    tuple(
+                        jax.ShapeDtypeStruct(shape, dtype)  # pyright: ignore[reportUnknownArgumentType]
+                        for shape, dtype in zip(shapes, dtypes, strict=True)
+                    ),
+                    *args,
+                    **kwargs,
                 ),
-                *args,
-                **kwargs,
-            ),
-        )
+            )
 
     else:
         # Eager backends