add xfail with request fixture

Author: Enderdead

src/array_api_extra/_delegation.py

@@ -19,6 +19,7 @@
 from ._lib._utils._typing import Array, DType
 
 __all__ = [
+    "atleast_nd",
     "cov",
     "expand_dims",
     "isclose",
@@ -29,6 +30,55 @@
 ]
 
 
+def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
+    """
+    Recursively expand the dimension of an array to at least `ndim`.
+
+    Parameters
+    ----------
+    x : array
+        Input array.
+    ndim : int
+        The minimum number of dimensions for the result.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `x`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array with ``res.ndim`` >= `ndim`.
+        If ``x.ndim`` >= `ndim`, `x` is returned.
+        If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes
+        until ``res.ndim`` equals `ndim`.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+    >>> x = xp.asarray([1])
+    >>> xpx.atleast_nd(x, ndim=3, xp=xp)
+    Array([[[1]]], dtype=array_api_strict.int64)
+
+    >>> x = xp.asarray([[[1, 2],
+    ...                  [3, 4]]])
+    >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x
+    True
+    """
+    if xp is None:
+        xp = array_namespace(x)
+
+    if 1 <= ndim <= 3 and (
+        is_numpy_namespace(xp)
+        or is_jax_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_torch_namespace(xp)
+    ):
+        return getattr(xp, f"atleast_{ndim}d")(x)
+
+    return _funcs.atleast_nd(x, ndim=ndim, xp=xp)
+
+
 def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
     """
     Estimate a covariance matrix.
@@ -197,55 +247,6 @@ def expand_dims(
     return _funcs.expand_dims(a, axis=axis, xp=xp)
 
 
-def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array:
-    """
-    Recursively expand the dimension of an array to at least `ndim`.
-
-    Parameters
-    ----------
-    x : array
-        Input array.
-    ndim : int
-        The minimum number of dimensions for the result.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `x`. Default: infer.
-
-    Returns
-    -------
-    array
-        An array with ``res.ndim`` >= `ndim`.
-        If ``x.ndim`` >= `ndim`, `x` is returned.
-        If ``x.ndim`` < `ndim`, `x` is expanded by prepending new axes
-        until ``res.ndim`` equals `ndim`.
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-    >>> x = xp.asarray([1])
-    >>> xpx.atleast_nd(x, ndim=3, xp=xp)
-    Array([[[1]]], dtype=array_api_strict.int64)
-
-    >>> x = xp.asarray([[[1, 2],
-    ...                  [3, 4]]])
-    >>> xpx.atleast_nd(x, ndim=1, xp=xp) is x
-    True
-    """
-    if xp is None:
-        xp = array_namespace(x)
-
-    if 1 <= ndim <= 3 and (
-        is_numpy_namespace(xp)
-        or is_jax_namespace(xp)
-        or is_dask_namespace(xp)
-        or is_cupy_namespace(xp)
-        or is_torch_namespace(xp)
-    ):
-        return getattr(xp, f"atleast_{ndim}d")(x)
-
-    return _funcs.atleast_nd(x, ndim=ndim, xp=xp)
-
-
 def isclose(
     a: Array | complex,
     b: Array | complex,

tests/test_funcs.py

@@ -33,7 +33,7 @@
     sinc,
 )
 from array_api_extra._lib._backends import NUMPY_VERSION, Backend
-from array_api_extra._lib._testing import xp_assert_close, xp_assert_equal
+from array_api_extra._lib._testing import xfail, xp_assert_close, xp_assert_equal
 from array_api_extra._lib._utils._compat import device as get_device
 from array_api_extra._lib._utils._compat import is_jax_namespace
 from array_api_extra._lib._utils._helpers import eager_shape, ndindex
@@ -1263,6 +1263,7 @@ def test_assume_unique(self, xp: ModuleType):
     @pytest.mark.parametrize("shape2", [(), (1,), (1, 1)])
     def test_shapes(
         self,
+        request: pytest.FixtureRequest,
         assume_unique: bool,
         shape1: tuple[int, ...],
         shape2: tuple[int, ...],
@@ -1272,22 +1273,24 @@ def test_shapes(
         x2 = xp.zeros(shape2)
 
         if is_jax_namespace(xp) and assume_unique and shape1 != (1,):
-            pytest.xfail(reason="jax#32335 fixed with jax>=0.8.0")
+            xfail(request=request, reason="jax#32335 fixed with jax>=0.8.0")
 
         actual = setdiff1d(x1, x2, assume_unique=assume_unique)
         xp_assert_equal(actual, xp.empty((0,)))
 
     @assume_unique
     @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp")
-    def test_python_scalar(self, xp: ModuleType, assume_unique: bool):
+    def test_python_scalar(
+        self, request: pytest.FixtureRequest, xp: ModuleType, assume_unique: bool
+    ):
         # Test no dtype promotion to xp.asarray(x2); use x1.dtype
         x1 = xp.asarray([3, 1, 2], dtype=xp.int16)
         x2 = 3
         actual = setdiff1d(x1, x2, assume_unique=assume_unique)
         xp_assert_equal(actual, xp.asarray([1, 2], dtype=xp.int16))
 
         if is_jax_namespace(xp) and assume_unique:
-            pytest.xfail(reason="jax#32335 fixed with jax>=0.8.0")
+            xfail(request=request, reason="jax#32335 fixed with jax>=0.8.0")
 
         actual = setdiff1d(x2, x1, assume_unique=assume_unique)
         xp_assert_equal(actual, xp.asarray([], dtype=xp.int16))