TST: cov: add another test

Author: lucascolley

src/array_api_extra/_funcs.py

@@ -49,28 +49,28 @@ def atleast_nd(x: Array, *, ndim: int, xp: ModuleType) -> Array:
     return x
 
 
-def cov(x: Array, *, xp: ModuleType) -> Array:
+def cov(m: Array, *, xp: ModuleType) -> Array:
     """..."""
-    x = xp.asarray(x, copy=True)
+    m = xp.asarray(m, copy=True)
     dtype = (
-        xp.float64 if xp.isdtype(x.dtype, "integral") else xp.result_type(x, xp.float64)
+        xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64)
     )
 
-    x = atleast_nd(x, ndim=2, xp=xp)
-    x = xp.astype(x, dtype)
+    m = atleast_nd(m, ndim=2, xp=xp)
+    m = xp.astype(m, dtype)
 
-    avg = mean(x, axis=1, xp=xp)
-    fact = x.shape[1] - 1
+    avg = mean(m, axis=1, xp=xp)
+    fact = m.shape[1] - 1
 
     if fact <= 0:
         warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2)
         fact = 0.0
 
-    x -= avg[:, None]
-    y_transpose = x.T
-    if xp.isdtype(y_transpose.dtype, "complex floating"):
-        y_transpose = xp.conj(y_transpose)
-    c = x @ y_transpose
+    m -= avg[:, None]
+    m_transpose = m.T
+    if xp.isdtype(m_transpose.dtype, "complex floating"):
+        m_transpose = xp.conj(m_transpose)
+    c = m @ m_transpose
     c /= fact
     axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
     return xp.squeeze(c, axis=axes)

tests/test_funcs.py

@@ -95,3 +95,12 @@ def test_empty(self):
                 cov(xp.reshape(xp.asarray([]), (2, 0)), xp=xp),
                 xp.asarray([[xp.nan, xp.nan], [xp.nan, xp.nan]]),
             )
+
+    def test_combination(self):
+        x = xp.asarray([-2.1, -1, 4.3])
+        y = xp.asarray([3, 1.1, 0.12])
+        X = xp.stack((x, y), axis=0)
+        desired = xp.asarray([[11.71, -4.286], [-4.286, 2.144133]])
+        assert_allclose(cov(X, xp=xp), desired, rtol=1e-6)
+        assert_allclose(cov(x, xp=xp), xp.asarray(11.71))
+        assert_allclose(cov(y, xp=xp), xp.asarray(2.144133), rtol=1e-6)