ENH: cov delegation (#451)

Co-authored-by: Adrián García Pitarch <[email protected]>
Co-authored-by: Lucas Colley <[email protected]>

Author: 3 people

src/array_api_extra/__init__.py

@@ -2,6 +2,7 @@
 
 from ._delegation import (
     atleast_nd,
+    cov,
     expand_dims,
     isclose,
     nan_to_num,
@@ -13,7 +14,6 @@
 from ._lib._funcs import (
     apply_where,
     broadcast_shapes,
-    cov,
     create_diagonal,
     default_dtype,
     kron,

src/array_api_extra/_delegation.py

@@ -18,7 +18,95 @@
 from ._lib._utils._helpers import asarrays
 from ._lib._utils._typing import Array, DType
 
-__all__ = ["expand_dims", "isclose", "nan_to_num", "one_hot", "pad", "sinc"]
+__all__ = [
+    "cov",
+    "expand_dims",
+    "isclose",
+    "nan_to_num",
+    "one_hot",
+    "pad",
+    "sinc",
+]
+
+
+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 expand_dims(

src/array_api_extra/_lib/_funcs.py

@@ -248,73 +248,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)

tests/test_funcs.py

@@ -419,46 +419,55 @@ def test_none(self, args: tuple[tuple[float | None, ...], ...]):
 class TestCov:
     def test_basic(self, xp: ModuleType):
         xp_assert_close(
-            cov(xp.asarray([[0, 2], [1, 1], [2, 0]]).T),
+            cov(xp.asarray([[0, 2], [1, 1], [2, 0]], dtype=xp.float64).T),
             xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64),
         )
 
     def test_complex(self, xp: ModuleType):
-        actual = cov(xp.asarray([[1, 2, 3], [1j, 2j, 3j]]))
+        actual = cov(xp.asarray([[1, 2, 3], [1j, 2j, 3j]], dtype=xp.complex128))
         expect = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]], dtype=xp.complex128)
         xp_assert_close(actual, expect)
 
+    @pytest.mark.xfail_xp_backend(Backend.JAX, reason="jax#32296")
     @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="sparse#877")
     def test_empty(self, xp: ModuleType):
         with warnings.catch_warnings(record=True):
             warnings.simplefilter("always", RuntimeWarning)
-            xp_assert_equal(cov(xp.asarray([])), xp.asarray(xp.nan, dtype=xp.float64))
+            warnings.simplefilter("always", UserWarning)
             xp_assert_equal(
-                cov(xp.reshape(xp.asarray([]), (0, 2))),
+                cov(xp.asarray([], dtype=xp.float64)),
+                xp.asarray(xp.nan, dtype=xp.float64),
+            )
+            xp_assert_equal(
+                cov(xp.reshape(xp.asarray([], dtype=xp.float64), (0, 2))),
                 xp.reshape(xp.asarray([], dtype=xp.float64), (0, 0)),
             )
             xp_assert_equal(
-                cov(xp.reshape(xp.asarray([]), (2, 0))),
+                cov(xp.reshape(xp.asarray([], dtype=xp.float64), (2, 0))),
                 xp.asarray([[xp.nan, xp.nan], [xp.nan, xp.nan]], dtype=xp.float64),
             )
 
     def test_combination(self, xp: ModuleType):
-        x = xp.asarray([-2.1, -1, 4.3])
-        y = xp.asarray([3, 1.1, 0.12])
+        x = xp.asarray([-2.1, -1, 4.3], dtype=xp.float64)
+        y = xp.asarray([3, 1.1, 0.12], dtype=xp.float64)
         X = xp.stack((x, y), axis=0)
         desired = xp.asarray([[11.71, -4.286], [-4.286, 2.144133]], dtype=xp.float64)
         xp_assert_close(cov(X), desired, rtol=1e-6)
         xp_assert_close(cov(x), xp.asarray(11.71, dtype=xp.float64))
         xp_assert_close(cov(y), xp.asarray(2.144133, dtype=xp.float64), rtol=1e-6)
 
+    @pytest.mark.xfail_xp_backend(Backend.TORCH, reason="array-api-extra#455")
     def test_device(self, xp: ModuleType, device: Device):
         x = xp.asarray([1, 2, 3], device=device)
         assert get_device(cov(x)) == device
 
     @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp")
     def test_xp(self, xp: ModuleType):
         xp_assert_close(
-            cov(xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]]).T, xp=xp),
+            cov(
+                xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]], dtype=xp.float64).T,
+                xp=xp,
+            ),
             xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64),
         )