ENH: cov delegation

Author: Adrián García Pitarch

src/array_api_extra/__init__.py

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

src/array_api_extra/_delegation.py

@@ -18,7 +18,87 @@
 from ._lib._utils._helpers import asarrays
 from ._lib._utils._typing import Array, DType
 
-__all__ = ["isclose", "nan_to_num", "one_hot", "pad"]
+__all__ = ["cov", "isclose", "nan_to_num", "one_hot", "pad"]
+
+
+def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
+    """
+    Estimate a covariance matrix.
+
+    Covariance indicates the level to which two variables vary together.
+    If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`,
+    then the covariance matrix element :math:`C_{ij}` is the covariance of
+    :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
+    of :math:`x_i`.
+
+    This provides a subset of the functionality of ``numpy.cov``.
+
+    Parameters
+    ----------
+    m : array
+        A 1-D or 2-D array containing multiple variables and observations.
+        Each row of `m` represents a variable, and each column a single
+        observation of all those variables.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `m`. Default: infer.
+
+    Returns
+    -------
+    array
+        The covariance matrix of the variables.
+
+    Examples
+    --------
+    >>> import array_api_strict as xp
+    >>> import array_api_extra as xpx
+
+    Consider two variables, :math:`x_0` and :math:`x_1`, which
+    correlate perfectly, but in opposite directions:
+
+    >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T
+    >>> x
+    Array([[0, 1, 2],
+           [2, 1, 0]], dtype=array_api_strict.int64)
+
+    Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance
+    matrix shows this clearly:
+
+    >>> xpx.cov(x, xp=xp)
+    Array([[ 1., -1.],
+           [-1.,  1.]], dtype=array_api_strict.float64)
+
+    Note that element :math:`C_{0,1}`, which shows the correlation between
+    :math:`x_0` and :math:`x_1`, is negative.
+
+    Further, note how `x` and `y` are combined:
+
+    >>> x = xp.asarray([-2.1, -1,  4.3])
+    >>> y = xp.asarray([3,  1.1,  0.12])
+    >>> X = xp.stack((x, y), axis=0)
+    >>> xpx.cov(X, xp=xp)
+    Array([[11.71      , -4.286     ],
+           [-4.286     ,  2.14413333]], dtype=array_api_strict.float64)
+
+    >>> xpx.cov(x, xp=xp)
+    Array(11.71, dtype=array_api_strict.float64)
+
+    >>> xpx.cov(y, xp=xp)
+    Array(2.14413333, dtype=array_api_strict.float64)
+    """
+
+    if xp is None:
+        xp = array_namespace(m)
+
+    if (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_torch_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_jax_namespace(xp)
+    ):
+        return xp.cov(m)
+
+    return _funcs.cov(m, xp=xp)
 
 
 def isclose(

src/array_api_extra/_lib/_funcs.py

@@ -281,73 +281,8 @@ def broadcast_shapes(*shapes: tuple[float | None, ...]) -> tuple[int | None, ...
     return tuple(out)
 
 
-def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array:
-    """
-    Estimate a covariance matrix.
-
-    Covariance indicates the level to which two variables vary together.
-    If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`,
-    then the covariance matrix element :math:`C_{ij}` is the covariance of
-    :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
-    of :math:`x_i`.
-
-    This provides a subset of the functionality of ``numpy.cov``.
-
-    Parameters
-    ----------
-    m : array
-        A 1-D or 2-D array containing multiple variables and observations.
-        Each row of `m` represents a variable, and each column a single
-        observation of all those variables.
-    xp : array_namespace, optional
-        The standard-compatible namespace for `m`. Default: infer.
-
-    Returns
-    -------
-    array
-        The covariance matrix of the variables.
-
-    Examples
-    --------
-    >>> import array_api_strict as xp
-    >>> import array_api_extra as xpx
-
-    Consider two variables, :math:`x_0` and :math:`x_1`, which
-    correlate perfectly, but in opposite directions:
-
-    >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T
-    >>> x
-    Array([[0, 1, 2],
-           [2, 1, 0]], dtype=array_api_strict.int64)
-
-    Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance
-    matrix shows this clearly:
-
-    >>> xpx.cov(x, xp=xp)
-    Array([[ 1., -1.],
-           [-1.,  1.]], dtype=array_api_strict.float64)
-
-    Note that element :math:`C_{0,1}`, which shows the correlation between
-    :math:`x_0` and :math:`x_1`, is negative.
-
-    Further, note how `x` and `y` are combined:
-
-    >>> x = xp.asarray([-2.1, -1,  4.3])
-    >>> y = xp.asarray([3,  1.1,  0.12])
-    >>> X = xp.stack((x, y), axis=0)
-    >>> xpx.cov(X, xp=xp)
-    Array([[11.71      , -4.286     ],
-           [-4.286     ,  2.14413333]], dtype=array_api_strict.float64)
-
-    >>> xpx.cov(x, xp=xp)
-    Array(11.71, dtype=array_api_strict.float64)
-
-    >>> xpx.cov(y, xp=xp)
-    Array(2.14413333, dtype=array_api_strict.float64)
-    """
-    if xp is None:
-        xp = array_namespace(m)
-
+def cov(m: Array, /, *, xp: ModuleType) -> Array:  # numpydoc ignore=PR01,RT01
+    """See docstring in array_api_extra._delegation."""
     m = xp.asarray(m, copy=True)
     dtype = (
         xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64)