ENH: validate weights shape in cov

Author: bruAristimunha

src/array_api_extra/_delegation.py

@@ -222,6 +222,22 @@ def cov(
     if m.ndim >= 2 and axis not in (-1, m.ndim - 1):
         m = xp.moveaxis(m, axis, -1)
 
+    # Validate weight shapes (eager metadata, lazy-safe). Value-based
+    # checks (non-negative, integer dtype) are intentionally skipped so
+    # lazy backends don't trigger compute -- same tradeoff as dask.cov.
+    n_obs = m.shape[-1]
+    for name, w in (("fweights", fweights), ("aweights", aweights)):
+        if w is None:
+            continue
+        if w.ndim != 1:
+            msg = f"`{name}` must be 1-D, got ndim={w.ndim}"
+            raise ValueError(msg)
+        if w.shape[0] != n_obs:
+            msg = (
+                f"`{name}` has length {w.shape[0]} but `m` has {n_obs} observations"
+            )
+            raise ValueError(msg)
+
     # `numpy.cov` (and cupy/dask/jax) require integer `ddof`; `torch.cov`
     # requires integer `correction`. For non-integer-valued `correction`,
     # fall through to the generic implementation.

tests/test_funcs.py

@@ -723,6 +723,19 @@ def test_axis_out_of_bounds(self, xp: ModuleType):
         with pytest.raises(IndexError):
             _ = cov(m, axis=5)
 
+    def test_weights_shape_validation(self, xp: ModuleType):
+        m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        # Wrong length.
+        with pytest.raises(ValueError, match="`fweights` has length"):
+            _ = cov(m, fweights=xp.asarray([1, 2]))
+        with pytest.raises(ValueError, match="`aweights` has length"):
+            _ = cov(m, aweights=xp.asarray([0.1, 0.2]))
+        # Wrong ndim.
+        with pytest.raises(ValueError, match="`fweights` must be 1-D"):
+            _ = cov(m, fweights=xp.asarray([[1, 2, 3]]))
+        with pytest.raises(ValueError, match="`aweights` must be 1-D"):
+            _ = cov(m, aweights=xp.asarray([[0.1, 0.2, 0.3]]))
+
 
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False)
 class TestOneHot: