MNT: move weights validation to generic cov

Author: bruAristimunha

src/array_api_extra/_delegation.py

@@ -222,20 +222,6 @@ 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.

src/array_api_extra/_lib/_funcs.py

@@ -299,6 +299,23 @@ def cov(
     m = atleast_nd(m, ndim=2, xp=xp)
     m = xp.astype(m, dtype)
 
+    # Validate weight shapes (eager metadata, lazy-safe). Native backends
+    # validate themselves; this covers the generic path (array-api-strict,
+    # sparse, and the dask+weights fallback where the native check is
+    # bypassed to preserve laziness).
+    n_obs = m.shape[-1]
+    for name, w_in in (("fweights", fweights), ("aweights", aweights)):
+        if w_in is None:
+            continue
+        if w_in.ndim != 1:
+            msg = f"`{name}` must be 1-D, got ndim={w_in.ndim}"
+            raise ValueError(msg)
+        if w_in.shape[0] != n_obs:
+            msg = (
+                f"`{name}` has length {w_in.shape[0]} but `m` has {n_obs} observations"
+            )
+            raise ValueError(msg)
+
     fw = None
     if fweights is not None:
         fw = xp.astype(xp.asarray(fweights), dtype)

tests/test_funcs.py

@@ -723,18 +723,6 @@ 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)