adress PR comments

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -15,7 +15,7 @@
     is_torch_namespace,
 )
 from ._lib._utils._compat import device as get_device
-from ._lib._utils._helpers import asarrays
+from ._lib._utils._helpers import asarrays, eager_shape
 from ._lib._utils._typing import Array, DType
 
 __all__ = ["isclose", "nan_to_num", "one_hot", "pad"]
@@ -364,10 +364,7 @@ def partition(
         raise TypeError(msg)
     if axis is None:
         return partition(xp.reshape(a, -1), kth, axis=0, xp=xp)
-    size = a.shape[axis]
-    if size is None:
-        msg = "Array dimensions must be known"
-        raise ValueError(msg)
+    (size,) = eager_shape(a, axis)
     if not (0 <= kth < size):
         msg = f"kth(={kth}) out of bounds [0 {size})"
         raise ValueError(msg)
@@ -445,10 +442,7 @@ def argpartition(
         raise TypeError(msg)
     if axis is None:
         return partition(xp.reshape(a, -1), kth, axis=0, xp=xp)
-    size = a.shape[axis]
-    if size is None:
-        msg = "Array dimensions must be known"
-        raise ValueError(msg)
+    (size,) = eager_shape(a, axis)
     if not (0 <= kth < size):
         msg = f"kth(={kth}) out of bounds [0 {size})"
         raise ValueError(msg)

src/array_api_extra/_lib/_utils/_helpers.py

@@ -250,22 +250,31 @@ def ndindex(*x: int) -> Generator[tuple[int, ...]]:
             yield *i, j
 
 
-def eager_shape(x: Array, /) -> tuple[int, ...]:
+def eager_shape(x: Array, /, axis: int | None = None) -> tuple[int, ...]:
     """
     Return shape of an array. Raise if shape is not fully defined.
 
     Parameters
     ----------
     x : Array
         Input array.
+    axis : int, optional
+        If provided, only returns the tuple (shape[axis],).
 
     Returns
     -------
     tuple[int, ...]
         Shape of the array.
     """
     shape = x.shape
-    # Dask arrays uses non-standard NaN instead of None
+    if axis is not None:
+        s = shape[axis]
+        # Dask arrays uses non-standard NaN instead of None
+        if s is None or math.isnan(s):
+            msg = f"Unsupported lazy shape for axis {axis}"
+            raise TypeError(msg)
+        return (s,)
+
     if any(s is None or math.isnan(s) for s in shape):
         msg = "Unsupported lazy shape"
         raise TypeError(msg)

tests/test_helpers.py

@@ -182,11 +182,14 @@ def test_eager_shape(xp: ModuleType, library: Backend):
     # Lazy arrays, like Dask, have an eager shape until you slice them with
     # a lazy boolean mask
     assert eager_shape(a) == a.shape == (3,)
+    assert eager_shape(a, axis=0) == a.shape == (3,)
 
     b = a[a > 2]
     if library is Backend.DASK:
         with pytest.raises(TypeError, match="Unsupported lazy shape"):
             _ = eager_shape(b)
+        with pytest.raises(TypeError, match="Unsupported lazy shape"):
+            _ = eager_shape(b, axis=0)
     # FIXME can't test use case for None in the shape until we add support for
     # other lazy backends
     else: