Revert bespoke exception

Author: crusaderky

docs/api-reference.md

@@ -7,7 +7,6 @@
     :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 UnknownShapeError, apply_numpy_func
+from ._lib._apply import apply_numpy_func
 from ._lib._funcs import (
     at,
     atleast_nd,
@@ -18,7 +18,6 @@
 
 # pylint: disable=duplicate-code
 __all__ = [
-    "UnknownShapeError",
     "__version__",
     "apply_numpy_func",
     "at",

src/array_api_extra/_lib/_apply.py

@@ -31,14 +31,6 @@ 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],
@@ -148,15 +140,14 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
 
     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.
+    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`.
+    RuntimeError
+        When `xp=sparse` and auto-densification is disabled.
+    Exception (backend-specific)
+        When the backend disallows implicit device to host transfers and the input
+        arrays are on a device, e.g. on GPU.
 
     See Also
     --------
@@ -241,15 +232,8 @@ def apply_numpy_func(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
         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
+            # Raises jax.errors.TracerArrayConversionError if we're inside jax.jit.
+            out = wrapped(*args, **kwargs)
 
         else:
             out = cast(