BUG: cov: preserve singleton batch dimensions

src/array_api_extra/_agnostic/_statistical.py

@@ -106,7 +106,8 @@ def cov(
     if xp.isdtype(m_cT.dtype, "complex floating"):
         m_cT = xp.conj(m_cT)
     c = m_w @ m_cT / fact
-    axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
+    # Match NumPy's scalar variance for one variable while retaining batch axes.
+    axes = tuple(axis for axis in (-2, -1) if c.shape[axis] == 1)
     return xp.squeeze(c, axis=axes)
 
 

src/array_api_extra/_statistical.py

@@ -73,7 +73,8 @@ def cov(
     -------
     array
         An array having shape ``(..., N, N)`` whose innermost two dimensions represent
-        the covariance matrix of the variables.
+        the covariance matrix of the variables. For a single variable, each covariance
+        matrix is returned as a scalar variance, retaining any batch dimensions.
 
     Notes
     -----

tests/main/test_statistical.py

@@ -93,14 +93,16 @@ def test_xp(self, xp: ArrayNamespace):
             xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64),
         )
 
-    def test_batch(self, xp: ArrayNamespace):
+    @pytest.mark.parametrize("batch_shape", [(3, 4), (1,), (2, 1), (1, 2, 1)])
+    @pytest.mark.parametrize("n_var", [1, 3])
+    def test_batch(self, batch_shape: tuple[int, ...], n_var: int, xp: ArrayNamespace):
         rng = np.random.default_rng(8847643423)
-        batch_shape = (3, 4)
-        n_var, n_obs = 3, 20
+        n_obs = 20
         m = rng.random((*batch_shape, n_var, n_obs))
         res = cov(xp.asarray(m))
         ref_list = [np.cov(m_) for m_ in np.reshape(m, (-1, n_var, n_obs))]
-        ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
+        output_shape = batch_shape if n_var == 1 else (*batch_shape, n_var, n_var)
+        ref = np.reshape(np.stack(ref_list), output_shape)
         assert_close(res, xp.asarray(ref))
 
     @pytest.mark.parametrize("bias", [True, False, 0, 1])