Merge branch 'main' into add-cov-delegation

Author: lucascolley

.all-contributorsrc

@@ -270,6 +270,16 @@
       "contributions": [
         "ideas"
       ]
+    },
+    {
+      "login": "Enderdead",
+      "name": "Gauthier-Clerc François",
+      "avatar_url": "https://avatars.githubusercontent.com/u/22777836?v=4",
+      "profile": "https://github.com/Enderdead",
+      "contributions": [
+        "code",
+        "test"
+      ]
     }
   ]
 }

CONTRIBUTORS.md

@@ -40,6 +40,7 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/pearu"><img src="https://avatars.githubusercontent.com/u/402156?v=4?s=100" width="100px;" alt="Pearu Peterson"/><br /><sub><b>Pearu Peterson</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Apearu" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://paddyroddy.github.io/"><img src="https://avatars.githubusercontent.com/u/15052188?v=4?s=100" width="100px;" alt="Patrick J. Roddy"/><br /><sub><b>Patrick J. Roddy</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Documentation">📖</a> <a href="#example-paddyroddy" title="Examples">💡</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://amacati.github.io/"><img src="https://avatars.githubusercontent.com/u/57562633?v=4?s=100" width="100px;" alt="Martin Schuck"/><br /><sub><b>Martin Schuck</b></sub></a><br /><a href="#ideas-amacati" title="Ideas, Planning, & Feedback">🤔</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/Enderdead"><img src="https://avatars.githubusercontent.com/u/22777836?v=4?s=100" width="100px;" alt="Gauthier-Clerc François"/><br /><sub><b>Gauthier-Clerc François</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=Enderdead" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=Enderdead" title="Tests">⚠️</a></td>
     </tr>
   </tbody>
 </table>

README.md

@@ -95,6 +95,7 @@ This project exists thanks to the following contributors
       <td align="center" valign="top" width="14.28%"><a href="https://github.com/pearu"><img src="https://avatars.githubusercontent.com/u/402156?v=4?s=100" width="100px;" alt="Pearu Peterson"/><br /><sub><b>Pearu Peterson</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/pulls?q=is%3Apr+reviewed-by%3Apearu" title="Reviewed Pull Requests">👀</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://paddyroddy.github.io/"><img src="https://avatars.githubusercontent.com/u/15052188?v=4?s=100" width="100px;" alt="Patrick J. Roddy"/><br /><sub><b>Patrick J. Roddy</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Documentation">📖</a> <a href="#example-paddyroddy" title="Examples">💡</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=paddyroddy" title="Tests">⚠️</a></td>
       <td align="center" valign="top" width="14.28%"><a href="https://amacati.github.io/"><img src="https://avatars.githubusercontent.com/u/57562633?v=4?s=100" width="100px;" alt="Martin Schuck"/><br /><sub><b>Martin Schuck</b></sub></a><br /><a href="#ideas-amacati" title="Ideas, Planning, & Feedback">🤔</a></td>
+      <td align="center" valign="top" width="14.28%"><a href="https://github.com/Enderdead"><img src="https://avatars.githubusercontent.com/u/22777836?v=4?s=100" width="100px;" alt="Gauthier-Clerc François"/><br /><sub><b>Gauthier-Clerc François</b></sub></a><br /><a href="https://github.com/data-apis/array-api-extra/commits?author=Enderdead" title="Code">💻</a> <a href="https://github.com/data-apis/array-api-extra/commits?author=Enderdead" title="Tests">⚠️</a></td>
     </tr>
   </tbody>
 </table>

pixi.lock

@@ -67,7 +67,7 @@ hypothesis = ">=6.138.3"
 dask-core = ">=2025.9.1" # No distributed, tornado, etc.
 dprint = ">=0.50.0,<0.51"
 lefthook = ">=1.13.1,<2"
-ruff = ">=0.12.12,<0.13"
+ruff = ">=0.13.1,<0.14"
 typos = ">=1.36.3,<2"
 actionlint = ">=1.7.7,<2"
 blacken-docs = ">=1.20.0,<2"

pyproject.toml

@@ -1,6 +1,7 @@
 """Extra array functions built on top of the array API standard."""
 
 from ._delegation import (
+    atleast_nd,
     cov,
     expand_dims,
     isclose,
@@ -12,7 +13,6 @@
 from ._lib._at import at
 from ._lib._funcs import (
     apply_where,
-    atleast_nd,
     broadcast_shapes,
     create_diagonal,
     default_dtype,

src/array_api_extra/__init__.py

@@ -197,6 +197,55 @@ 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,

src/array_api_extra/_delegation.py

@@ -175,42 +175,9 @@ def _apply_where(  # numpydoc ignore=PR01,RT01
     return at(out, cond).set(temp1)
 
 
-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)
+def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType) -> Array:
+    # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
 
     if x.ndim < ndim:
         x = xp.expand_dims(x, axis=0)