From 599b3c7acb6f0784b80ef693e6d2aa5717607a1a Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Fri, 17 Apr 2026 12:37:31 +0200
Subject: [PATCH 01/16] ENH: expose correction and weights parameters in cov
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Resolves #688. Adds `axis`, `correction`, `frequency_weights`, and
`weights` to `cov`, giving users control over the degrees-of-freedom
correction and the observation-axis / weighted variants that
`numpy.cov` and `torch.cov` already support.

Naming follows array-api conventions (`axis`, `correction`) rather
than numpy's (`rowvar`, `bias`, `ddof`); the docstring includes a
one-to-one mapping. The delegation moves observations to the last
axis via `xp.moveaxis`, collapsing `rowvar` out of the backend
dispatch — only `ddof` vs `correction` differs between branches.

Dask's native `cov` forces `.compute()` on a lazy scalar when any
weights are given, so weighted dask inputs fall through to the
generic implementation, which is fully lazy.
---
 src/array_api_extra/_delegation.py | 120 +++++++++++++++++++++++++----
 src/array_api_extra/_lib/_funcs.py |  66 ++++++++++++----
 tests/test_funcs.py                |  91 ++++++++++++++++++++++
 3 files changed, 246 insertions(+), 31 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index c5c16a37..150039d6 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -166,7 +166,16 @@ def broadcast_shapes(
     return _funcs.broadcast_shapes(*shapes)
 
 
-def cov(m: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
+def cov(
+    m: Array,
+    /,
+    *,
+    axis: int = -1,
+    correction: int | float = 1,
+    frequency_weights: Array | None = None,
+    weights: Array | None = None,
+    xp: ArrayNamespace | None = None,
+) -> Array:
     """
     Estimate a covariance matrix (or a stack of covariance matrices).
 
@@ -177,16 +186,37 @@ def cov(m: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
     :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
     of :math:`x_i`.
 
-    With the exception of supporting batch input, this provides a subset of
-    the functionality of ``numpy.cov``.
+    Extends ``numpy.cov`` with support for batch input and array-api
+    backends. Naming follows the array-api conventions used elsewhere in
+    this library (``axis``, ``correction``) rather than the numpy spellings
+    (``rowvar``, ``bias``, ``ddof``); see Notes for the mapping.
 
     Parameters
     ----------
     m : array
         An array of shape ``(..., N, M)`` whose innermost two dimensions
-        contain *M* observations of *N* variables. That is,
-        each row of `m` represents a variable, and each column a single
-        observation of all those variables.
+        contain *M* observations of *N* variables by default. The axis of
+        observations is controlled by `axis`.
+    axis : int, optional
+        Axis of `m` containing the observations. Default: ``-1`` (the last
+        axis), matching the array-api convention. Use ``axis=-2`` (or ``0``
+        for 2-D input) to treat each column as a variable, which
+        corresponds to ``rowvar=False`` in ``numpy.cov``.
+    correction : int or float, optional
+        Degrees of freedom correction: normalization divides by
+        ``N - correction`` (for unweighted input). Default: ``1``, which
+        gives the unbiased estimate (matches ``numpy.cov`` default of
+        ``bias=False``). Set to ``0`` for the biased estimate (``N``
+        normalization). Corresponds to ``ddof`` in ``numpy.cov`` and to
+        ``correction`` in ``numpy.var``/``std`` and ``torch.cov``.
+    frequency_weights : array, optional
+        1-D array of integer frequency weights: the number of times each
+        observation is repeated. Corresponds to ``fweights`` in
+        ``numpy.cov``/``torch.cov``.
+    weights : array, optional
+        1-D array of observation-vector weights (analytic weights). Larger
+        values mark more important observations. Corresponds to
+        ``aweights`` in ``numpy.cov``/``torch.cov``.
     xp : array_namespace, optional
         The standard-compatible namespace for `m`. Default: infer.
 
@@ -196,6 +226,23 @@ def cov(m: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
         An array having shape (..., N, N) whose innermost two dimensions represent
         the covariance matrix of the variables.
 
+    Notes
+    -----
+    Mapping from ``numpy.cov`` to this function::
+
+        numpy.cov(m, rowvar=True)           -> cov(m, axis=-1)   # default
+        numpy.cov(m, rowvar=False)          -> cov(m, axis=-2)
+        numpy.cov(m, bias=True)             -> cov(m, correction=0)
+        numpy.cov(m, ddof=k)                -> cov(m, correction=k)
+        numpy.cov(m, fweights=f)            -> cov(m, frequency_weights=f)
+        numpy.cov(m, aweights=a)            -> cov(m, weights=a)
+
+    Unlike ``numpy.cov``, a ``RuntimeWarning`` for non-positive effective
+    degrees of freedom is only emitted on the unweighted path. The
+    weighted path omits the check so that lazy backends (e.g. Dask) can
+    stay lazy end-to-end; choose ``correction`` and weights such that the
+    effective normalizer is positive.
+
     Examples
     --------
     >>> import array_api_strict as xp
@@ -249,16 +296,57 @@ def cov(m: Array, /, *, xp: ArrayNamespace | None = None) -> Array:
     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)
-    ) and m.ndim <= 2:
-        return xp.cov(m)
-
-    return _funcs.cov(m, xp=xp)
+    # Validate axis against m.ndim.
+    ndim = max(m.ndim, 1)
+    if not -ndim <= axis < ndim:
+        msg = f"axis {axis} is out of bounds for array of dimension {m.ndim}"
+        raise IndexError(msg)
+
+    # Normalize: observations on the last axis. After this, every backend
+    # sees the same convention and we never need to deal with `rowvar`.
+    if m.ndim >= 2 and axis not in (-1, m.ndim - 1):
+        m = xp.moveaxis(m, axis, -1)
+
+    # `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.
+    integer_correction = isinstance(correction, int) or correction.is_integer()
+    has_weights = frequency_weights is not None or weights is not None
+
+    if m.ndim <= 2 and integer_correction:
+        if is_torch_namespace(xp):
+            device = get_device(m)
+            fw = (
+                None
+                if frequency_weights is None
+                else xp.asarray(frequency_weights, device=device)
+            )
+            aw = None if weights is None else xp.asarray(weights, device=device)
+            return xp.cov(m, correction=int(correction), fweights=fw, aweights=aw)
+        # `dask.array.cov` forces `.compute()` whenever weights are given:
+        # its internal `if fact <= 0` check on a lazy 0-D scalar triggers
+        # materialization. Route to the generic impl, which is fully lazy
+        # because it only does sum/matmul and skips that scalar check.
+        if (
+            is_numpy_namespace(xp)
+            or is_cupy_namespace(xp)
+            or is_jax_namespace(xp)
+            or (is_dask_namespace(xp) and not has_weights)
+        ):
+            return xp.cov(
+                m,
+                ddof=int(correction),
+                fweights=frequency_weights,
+                aweights=weights,
+            )
+
+    return _funcs.cov(
+        m,
+        correction=correction,
+        frequency_weights=frequency_weights,
+        weights=weights,
+        xp=xp,
+    )
 
 
 def create_diagonal(
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 3f1ef511..e59b880d 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -260,9 +260,17 @@ def broadcast_shapes(  # numpydoc ignore=PR01,RT01
     return tuple(out)
 
 
-def cov(m: Array, /, *, xp: ArrayNamespace) -> Array:  # numpydoc ignore=PR01,RT01
+def cov(
+    m: Array,
+    /,
+    *,
+    correction: int | float = 1,
+    frequency_weights: Array | None = None,
+    weights: Array | None = None,
+    xp: ArrayNamespace,
+) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""
-    m = xp.asarray(m, copy=True)
+    m = xp.asarray(m)
     dtype = (
         xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64)
     )
@@ -270,21 +278,49 @@ def cov(m: Array, /, *, xp: ArrayNamespace) -> Array:  # numpydoc ignore=PR01,RT
     m = atleast_nd(m, ndim=2, xp=xp)
     m = xp.astype(m, dtype)
 
-    avg = xp.mean(m, axis=-1, keepdims=True)
+    device = _compat.device(m)
+    fw = (
+        None
+        if frequency_weights is None
+        else xp.astype(xp.asarray(frequency_weights, device=device), dtype)
+    )
+    aw = (
+        None
+        if weights is None
+        else xp.astype(xp.asarray(weights, device=device), dtype)
+    )
+    if fw is None and aw is None:
+        w = None
+    elif fw is None:
+        w = aw
+    elif aw is None:
+        w = fw
+    else:
+        w = fw * aw
 
     m_shape = eager_shape(m)
-    fact = m_shape[-1] - 1
-
-    if fact <= 0:
-        warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2)
-        fact = 0
-
-    m -= avg
-    m_transpose = xp.matrix_transpose(m)
-    if xp.isdtype(m_transpose.dtype, "complex floating"):
-        m_transpose = xp.conj(m_transpose)
-    c = xp.matmul(m, m_transpose)
-    c /= fact
+    if w is None:
+        avg = xp.mean(m, axis=-1, keepdims=True)
+        fact = m_shape[-1] - correction
+        if fact <= 0:
+            warnings.warn(
+                "Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2
+            )
+            fact = 0
+    else:
+        v1 = xp.sum(w, axis=-1)
+        avg = xp.sum(m * w, axis=-1, keepdims=True) / v1
+        if aw is None:
+            fact = v1 - correction
+        else:
+            fact = v1 - correction * xp.sum(w * aw, axis=-1) / v1
+
+    m_c = m - avg
+    m_w = m_c if w is None else m_c * w
+    m_cT = xp.matrix_transpose(m_c)
+    if xp.isdtype(m_cT.dtype, "complex floating"):
+        m_cT = xp.conj(m_cT)
+    c = xp.matmul(m_w, m_cT) / fact
     axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
     return xp.squeeze(c, axis=axes)
 
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index 8b639564..b5ff66e5 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -668,6 +668,97 @@ def test_batch(self, xp: ArrayNamespace):
         ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
         assert_close(res, xp.asarray(ref))
 
+    def test_correction(self, xp: ModuleType):
+        rng = np.random.default_rng(20260417)
+        m = rng.random((3, 20))
+        for correction in (0, 1, 2):
+            ref = np.cov(m, ddof=correction)
+            res = cov(xp.asarray(m), correction=correction)
+            xp_assert_close(res, xp.asarray(ref))
+
+    def test_correction_float(self, xp: ModuleType):
+        # Float correction: reference computed by hand (numpy.cov rejects
+        # non-integer ddof; our generic path supports it).
+        rng = np.random.default_rng(20260417)
+        m = rng.random((3, 20))
+        n = m.shape[-1]
+        centered = m - m.mean(axis=-1, keepdims=True)
+        ref = centered @ centered.T / (n - 1.5)
+        res = cov(xp.asarray(m), correction=1.5)
+        xp_assert_close(res, xp.asarray(ref))
+
+    def test_axis(self, xp: ModuleType):
+        rng = np.random.default_rng(20260417)
+        m = rng.random((20, 3))  # observations on axis 0
+        ref = np.cov(m, rowvar=False)
+        res = cov(xp.asarray(m), axis=0)
+        xp_assert_close(res, xp.asarray(ref))
+        res_neg = cov(xp.asarray(m), axis=-2)
+        xp_assert_close(res_neg, xp.asarray(ref))
+
+    def test_frequency_weights(self, xp: ModuleType):
+        rng = np.random.default_rng(20260417)
+        m = rng.random((3, 10))
+        fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64)
+        ref = np.cov(m, fweights=fw)
+        res = cov(xp.asarray(m), frequency_weights=xp.asarray(fw))
+        xp_assert_close(res, xp.asarray(ref))
+
+    def test_weights(self, xp: ModuleType):
+        rng = np.random.default_rng(20260417)
+        m = rng.random((3, 10))
+        aw = rng.random(10)
+        ref = np.cov(m, aweights=aw)
+        res = cov(xp.asarray(m), weights=xp.asarray(aw))
+        xp_assert_close(res, xp.asarray(ref))
+
+    def test_both_weights(self, xp: ModuleType):
+        rng = np.random.default_rng(20260417)
+        m = rng.random((3, 10))
+        fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64)
+        aw = rng.random(10)
+        for correction in (0, 1, 2):
+            ref = np.cov(m, ddof=correction, fweights=fw, aweights=aw)
+            res = cov(
+                xp.asarray(m),
+                correction=correction,
+                frequency_weights=xp.asarray(fw),
+                weights=xp.asarray(aw),
+            )
+            xp_assert_close(res, xp.asarray(ref))
+
+    def test_batch_with_weights(self, xp: ModuleType):
+        rng = np.random.default_rng(20260417)
+        batch_shape = (2, 3)
+        n_var, n_obs = 3, 15
+        m = rng.random((*batch_shape, n_var, n_obs))
+        aw = rng.random(n_obs)
+        res = cov(xp.asarray(m), weights=xp.asarray(aw))
+        ref_list = [np.cov(m_, aweights=aw) for m_ in np.reshape(m, (-1, n_var, n_obs))]
+        ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
+        xp_assert_close(res, xp.asarray(ref))
+
+    def test_axis_with_weights(self, xp: ModuleType):
+        # axis=-2 (observations on first of 2D) combined with weights:
+        # verifies that moveaxis and weight alignment cooperate.
+        rng = np.random.default_rng(20260417)
+        m = rng.random((15, 3))  # observations on axis 0
+        aw = rng.random(15)
+        fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1, 1, 1, 2, 1, 1], dtype=np.int64)
+        ref = np.cov(m, rowvar=False, fweights=fw, aweights=aw)
+        res = cov(
+            xp.asarray(m),
+            axis=-2,
+            frequency_weights=xp.asarray(fw),
+            weights=xp.asarray(aw),
+        )
+        xp_assert_close(res, xp.asarray(ref))
+
+    def test_axis_out_of_bounds(self, xp: ModuleType):
+        m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        with pytest.raises(IndexError):
+            _ = cov(m, axis=5)
+
 
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False)
 class TestOneHot:

From 8f454d3e24461655539e291d66fe0ac7657ee4a3 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 12:06:47 +0200
Subject: [PATCH 02/16] MNT: drop device= in cov weights

---
 src/array_api_extra/_delegation.py | 7 ++-----
 src/array_api_extra/_lib/_funcs.py | 5 ++---
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 150039d6..cf3c370d 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -315,13 +315,10 @@ def cov(
 
     if m.ndim <= 2 and integer_correction:
         if is_torch_namespace(xp):
-            device = get_device(m)
             fw = (
-                None
-                if frequency_weights is None
-                else xp.asarray(frequency_weights, device=device)
+                None if frequency_weights is None else xp.asarray(frequency_weights)
             )
-            aw = None if weights is None else xp.asarray(weights, device=device)
+            aw = None if weights is None else xp.asarray(weights)
             return xp.cov(m, correction=int(correction), fweights=fw, aweights=aw)
         # `dask.array.cov` forces `.compute()` whenever weights are given:
         # its internal `if fact <= 0` check on a lazy 0-D scalar triggers
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index e59b880d..8d7086db 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -278,16 +278,15 @@ def cov(
     m = atleast_nd(m, ndim=2, xp=xp)
     m = xp.astype(m, dtype)
 
-    device = _compat.device(m)
     fw = (
         None
         if frequency_weights is None
-        else xp.astype(xp.asarray(frequency_weights, device=device), dtype)
+        else xp.astype(xp.asarray(frequency_weights), dtype)
     )
     aw = (
         None
         if weights is None
-        else xp.astype(xp.asarray(weights, device=device), dtype)
+        else xp.astype(xp.asarray(weights), dtype)
     )
     if fw is None and aw is None:
         w = None

From aefcc41bef2de74fad1c0ebf126bdcce2dd414f8 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 12:11:49 +0200
Subject: [PATCH 03/16] STY: formatter

---
 src/array_api_extra/_delegation.py | 4 +---
 src/array_api_extra/_lib/_funcs.py | 6 +-----
 2 files changed, 2 insertions(+), 8 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index cf3c370d..d4a638e4 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -315,9 +315,7 @@ def cov(
 
     if m.ndim <= 2 and integer_correction:
         if is_torch_namespace(xp):
-            fw = (
-                None if frequency_weights is None else xp.asarray(frequency_weights)
-            )
+            fw = None if frequency_weights is None else xp.asarray(frequency_weights)
             aw = None if weights is None else xp.asarray(weights)
             return xp.cov(m, correction=int(correction), fweights=fw, aweights=aw)
         # `dask.array.cov` forces `.compute()` whenever weights are given:
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 8d7086db..bf23e95b 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -283,11 +283,7 @@ def cov(
         if frequency_weights is None
         else xp.astype(xp.asarray(frequency_weights), dtype)
     )
-    aw = (
-        None
-        if weights is None
-        else xp.astype(xp.asarray(weights), dtype)
-    )
+    aw = None if weights is None else xp.astype(xp.asarray(weights), dtype)
     if fw is None and aw is None:
         w = None
     elif fw is None:

From d3ad23e6a67248f55c9ad5a4471d2688dc4b2b48 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 12:20:17 +0200
Subject: [PATCH 04/16] TST: add bias tests from #691

---
 tests/test_funcs.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index b5ff66e5..92d6b919 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -668,6 +668,30 @@ def test_batch(self, xp: ArrayNamespace):
         ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
         assert_close(res, xp.asarray(ref))
 
+    @pytest.mark.parametrize("bias", [True, False, 0, 1])
+    def test_bias(self, xp: ModuleType, bias: bool):
+        # `bias` maps to `correction`: bias=True -> correction=0, bias=False -> 1.
+        x = np.array([-2.1, -1, 4.3])
+        y = np.array([3, 1.1, 0.12])
+        X = np.stack((x, y), axis=0)
+        ref = np.cov(X, bias=bias)
+        xp_assert_close(
+            cov(xp.asarray(X, dtype=xp.float64), correction=0 if bias else 1),
+            xp.asarray(ref, dtype=xp.float64),
+            rtol=1e-6,
+        )
+
+    @pytest.mark.parametrize("bias", [True, False, 0, 1])
+    def test_bias_batch(self, xp: ModuleType, bias: bool):
+        rng = np.random.default_rng(8847643423)
+        batch_shape = (3, 4)
+        n_var, n_obs = 3, 20
+        m = rng.random((*batch_shape, n_var, n_obs))
+        res = cov(xp.asarray(m), correction=0 if bias else 1)
+        ref_list = [np.cov(m_, bias=bias) for m_ in np.reshape(m, (-1, n_var, n_obs))]
+        ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
+        xp_assert_close(res, xp.asarray(ref))
+
     def test_correction(self, xp: ModuleType):
         rng = np.random.default_rng(20260417)
         m = rng.random((3, 20))

From 35a00896810d7e70bcb415351b4b45c346f47771 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 13:26:20 +0200
Subject: [PATCH 05/16] Update _funcs.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>
---
 src/array_api_extra/_lib/_funcs.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index bf23e95b..dbff8fa4 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -278,11 +278,9 @@ def cov(
     m = atleast_nd(m, ndim=2, xp=xp)
     m = xp.astype(m, dtype)
 
-    fw = (
-        None
-        if frequency_weights is None
-        else xp.astype(xp.asarray(frequency_weights), dtype)
-    )
+    fw = None
+    if frequency_weights is not None:
+        fw = xp.astype(xp.asarray(frequency_weights), dtype)
     aw = None if weights is None else xp.astype(xp.asarray(weights), dtype)
     if fw is None and aw is None:
         w = None

From 32fc1f71ac3db01d73e0e8f0df83f8d74c104653 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 13:39:04 +0200
Subject: [PATCH 06/16] Update _delegation.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>
---
 src/array_api_extra/_delegation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index d4a638e4..50c5cba7 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -209,11 +209,11 @@ def cov(
         ``bias=False``). Set to ``0`` for the biased estimate (``N``
         normalization). Corresponds to ``ddof`` in ``numpy.cov`` and to
         ``correction`` in ``numpy.var``/``std`` and ``torch.cov``.
-    frequency_weights : array, optional
+    fweights : array, optional
         1-D array of integer frequency weights: the number of times each
         observation is repeated. Corresponds to ``fweights`` in
         ``numpy.cov``/``torch.cov``.
-    weights : array, optional
+    aweights : array, optional
         1-D array of observation-vector weights (analytic weights). Larger
         values mark more important observations. Corresponds to
         ``aweights`` in ``numpy.cov``/``torch.cov``.

From 0ce74d8bdb84070f975cc93da6ad1be391ffa505 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 13:57:08 +0200
Subject: [PATCH 07/16] MNT: rename weights params to fweights/aweights

---
 src/array_api_extra/_delegation.py | 28 ++++++++++++++--------------
 src/array_api_extra/_lib/_funcs.py | 12 +++++++-----
 tests/test_funcs.py                | 14 +++++++-------
 3 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 50c5cba7..7665b4ed 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -172,8 +172,8 @@ def cov(
     *,
     axis: int = -1,
     correction: int | float = 1,
-    frequency_weights: Array | None = None,
-    weights: Array | None = None,
+    fweights: Array | None = None,
+    aweights: Array | None = None,
     xp: ArrayNamespace | None = None,
 ) -> Array:
     """
@@ -211,12 +211,12 @@ def cov(
         ``correction`` in ``numpy.var``/``std`` and ``torch.cov``.
     fweights : array, optional
         1-D array of integer frequency weights: the number of times each
-        observation is repeated. Corresponds to ``fweights`` in
+        observation is repeated. Same as ``fweights`` in
         ``numpy.cov``/``torch.cov``.
     aweights : array, optional
         1-D array of observation-vector weights (analytic weights). Larger
-        values mark more important observations. Corresponds to
-        ``aweights`` in ``numpy.cov``/``torch.cov``.
+        values mark more important observations. Same as ``aweights`` in
+        ``numpy.cov``/``torch.cov``.
     xp : array_namespace, optional
         The standard-compatible namespace for `m`. Default: infer.
 
@@ -234,8 +234,8 @@ def cov(
         numpy.cov(m, rowvar=False)          -> cov(m, axis=-2)
         numpy.cov(m, bias=True)             -> cov(m, correction=0)
         numpy.cov(m, ddof=k)                -> cov(m, correction=k)
-        numpy.cov(m, fweights=f)            -> cov(m, frequency_weights=f)
-        numpy.cov(m, aweights=a)            -> cov(m, weights=a)
+        numpy.cov(m, fweights=f)            -> cov(m, fweights=f)
+        numpy.cov(m, aweights=a)            -> cov(m, aweights=a)
 
     Unlike ``numpy.cov``, a ``RuntimeWarning`` for non-positive effective
     degrees of freedom is only emitted on the unweighted path. The
@@ -311,12 +311,12 @@ def cov(
     # requires integer `correction`. For non-integer-valued `correction`,
     # fall through to the generic implementation.
     integer_correction = isinstance(correction, int) or correction.is_integer()
-    has_weights = frequency_weights is not None or weights is not None
+    has_weights = fweights is not None or aweights is not None
 
     if m.ndim <= 2 and integer_correction:
         if is_torch_namespace(xp):
-            fw = None if frequency_weights is None else xp.asarray(frequency_weights)
-            aw = None if weights is None else xp.asarray(weights)
+            fw = None if fweights is None else xp.asarray(fweights)
+            aw = None if aweights is None else xp.asarray(aweights)
             return xp.cov(m, correction=int(correction), fweights=fw, aweights=aw)
         # `dask.array.cov` forces `.compute()` whenever weights are given:
         # its internal `if fact <= 0` check on a lazy 0-D scalar triggers
@@ -331,15 +331,15 @@ def cov(
             return xp.cov(
                 m,
                 ddof=int(correction),
-                fweights=frequency_weights,
-                aweights=weights,
+                fweights=fweights,
+                aweights=aweights,
             )
 
     return _funcs.cov(
         m,
         correction=correction,
-        frequency_weights=frequency_weights,
-        weights=weights,
+        fweights=fweights,
+        aweights=aweights,
         xp=xp,
     )
 
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index dbff8fa4..7d9ba611 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -265,8 +265,8 @@ def cov(
     /,
     *,
     correction: int | float = 1,
-    frequency_weights: Array | None = None,
-    weights: Array | None = None,
+    fweights: Array | None = None,
+    aweights: Array | None = None,
     xp: ArrayNamespace,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""
@@ -279,9 +279,11 @@ def cov(
     m = xp.astype(m, dtype)
 
     fw = None
-    if frequency_weights is not None:
-        fw = xp.astype(xp.asarray(frequency_weights), dtype)
-    aw = None if weights is None else xp.astype(xp.asarray(weights), dtype)
+    if fweights is not None:
+        fw = xp.astype(xp.asarray(fweights), dtype)
+    aw = None
+    if aweights is not None:
+        aw = xp.astype(xp.asarray(aweights), dtype)
     if fw is None and aw is None:
         w = None
     elif fw is None:
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index 92d6b919..c06555c1 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -725,7 +725,7 @@ def test_frequency_weights(self, xp: ModuleType):
         m = rng.random((3, 10))
         fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64)
         ref = np.cov(m, fweights=fw)
-        res = cov(xp.asarray(m), frequency_weights=xp.asarray(fw))
+        res = cov(xp.asarray(m), fweights=xp.asarray(fw))
         xp_assert_close(res, xp.asarray(ref))
 
     def test_weights(self, xp: ModuleType):
@@ -733,7 +733,7 @@ def test_weights(self, xp: ModuleType):
         m = rng.random((3, 10))
         aw = rng.random(10)
         ref = np.cov(m, aweights=aw)
-        res = cov(xp.asarray(m), weights=xp.asarray(aw))
+        res = cov(xp.asarray(m), aweights=xp.asarray(aw))
         xp_assert_close(res, xp.asarray(ref))
 
     def test_both_weights(self, xp: ModuleType):
@@ -746,8 +746,8 @@ def test_both_weights(self, xp: ModuleType):
             res = cov(
                 xp.asarray(m),
                 correction=correction,
-                frequency_weights=xp.asarray(fw),
-                weights=xp.asarray(aw),
+                fweights=xp.asarray(fw),
+                aweights=xp.asarray(aw),
             )
             xp_assert_close(res, xp.asarray(ref))
 
@@ -757,7 +757,7 @@ def test_batch_with_weights(self, xp: ModuleType):
         n_var, n_obs = 3, 15
         m = rng.random((*batch_shape, n_var, n_obs))
         aw = rng.random(n_obs)
-        res = cov(xp.asarray(m), weights=xp.asarray(aw))
+        res = cov(xp.asarray(m), aweights=xp.asarray(aw))
         ref_list = [np.cov(m_, aweights=aw) for m_ in np.reshape(m, (-1, n_var, n_obs))]
         ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
         xp_assert_close(res, xp.asarray(ref))
@@ -773,8 +773,8 @@ def test_axis_with_weights(self, xp: ModuleType):
         res = cov(
             xp.asarray(m),
             axis=-2,
-            frequency_weights=xp.asarray(fw),
-            weights=xp.asarray(aw),
+            fweights=xp.asarray(fw),
+            aweights=xp.asarray(aw),
         )
         xp_assert_close(res, xp.asarray(ref))
 

From b13185dfb733540966373fe38ed6f2ee8d0056a8 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 14:04:35 +0200
Subject: [PATCH 08/16] ENH: validate weights shape in cov

---
 src/array_api_extra/_delegation.py | 16 ++++++++++++++++
 tests/test_funcs.py                | 13 +++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 7665b4ed..506a897b 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -307,6 +307,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.
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index c06555c1..f2dea18e 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -783,6 +783,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:

From f859d209455febaaa98a294cfd97cdab587eb37d Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 15:35:17 +0200
Subject: [PATCH 09/16] MNT: address lucascolley review

---
 src/array_api_extra/_delegation.py | 6 ++----
 src/array_api_extra/_lib/_funcs.py | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 506a897b..3411a5e6 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -294,7 +294,7 @@ def cov(
     """
 
     if xp is None:
-        xp = array_namespace(m)
+        xp = array_namespace(m, fweights, aweights)
 
     # Validate axis against m.ndim.
     ndim = max(m.ndim, 1)
@@ -318,9 +318,7 @@ def cov(
             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"
-            )
+            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`
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 7d9ba611..90d7101c 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -315,7 +315,7 @@ def cov(
     m_cT = xp.matrix_transpose(m_c)
     if xp.isdtype(m_cT.dtype, "complex floating"):
         m_cT = xp.conj(m_cT)
-    c = xp.matmul(m_w, m_cT) / fact
+    c = (m_w @ m_cT) / fact
     axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
     return xp.squeeze(c, axis=axes)
 

From 29a400ae2f30e4af7e30e2d996603551af01c153 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 16:01:19 +0200
Subject: [PATCH 10/16] MNT: move weights validation to generic cov

---
 src/array_api_extra/_delegation.py | 14 --------------
 src/array_api_extra/_lib/_funcs.py | 17 +++++++++++++++++
 tests/test_funcs.py                | 12 ------------
 3 files changed, 17 insertions(+), 26 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 3411a5e6..ff380730 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -307,20 +307,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.
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 90d7101c..88af3d44 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -278,6 +278,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)
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index f2dea18e..a9f7ac68 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -783,18 +783,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)

From 75c6cfe584dc226d5f89a698a3aeb5d577cd03e9 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Mon, 20 Apr 2026 22:41:40 +0200
Subject: [PATCH 11/16] Update _funcs.py
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>
---
 src/array_api_extra/_lib/_funcs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 88af3d44..318e1b02 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -332,7 +332,7 @@ def cov(
     m_cT = xp.matrix_transpose(m_c)
     if xp.isdtype(m_cT.dtype, "complex floating"):
         m_cT = xp.conj(m_cT)
-    c = (m_w @ m_cT) / fact
+    c = m_w @ m_cT / fact
     axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
     return xp.squeeze(c, axis=axes)
 

From 1fdf3187f70b84b6c24e0a6373a6503eec1fdd1a Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Fri, 24 Apr 2026 15:13:19 +0200
Subject: [PATCH 12/16] DOC: explain non-integer correction use cases in cov

Addresses review feedback (kgryte, betatim) that the motivation for
allowing non-integer correction was not obvious from the docstring:
weighted unbiased correction and autocorrelated data both require
fractional values.
---
 src/array_api_extra/_delegation.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index ff380730..293c3ba5 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -209,6 +209,13 @@ def cov(
         ``bias=False``). Set to ``0`` for the biased estimate (``N``
         normalization). Corresponds to ``ddof`` in ``numpy.cov`` and to
         ``correction`` in ``numpy.var``/``std`` and ``torch.cov``.
+        Non-integer values are allowed for advanced use cases: the
+        unbiased correction for weighted observations depends on the
+        sum and dispersion of the weights and is generally not an
+        integer, and autocorrelated data may also require a fractional
+        correction. Non-integer ``correction`` routes through the
+        generic implementation because ``numpy.cov``'s ``ddof`` and
+        ``torch.cov``'s ``correction`` both require integers.
     fweights : array, optional
         1-D array of integer frequency weights: the number of times each
         observation is repeated. Same as ``fweights`` in

From 8d1989c84ff1b71a7f253c4db7f3cbee7b52778c Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Fri, 24 Apr 2026 15:13:30 +0200
Subject: [PATCH 13/16] TST: cover weight validation error paths in cov

Adds tests for the 1-D shape and length checks in the generic cov
path. Raises the diff coverage for this PR from 93.33% to 100%.
---
 tests/test_funcs.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index a9f7ac68..d832e2d0 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -783,6 +783,23 @@ def test_axis_out_of_bounds(self, xp: ModuleType):
         with pytest.raises(IndexError):
             _ = cov(m, axis=5)
 
+    def test_weights_wrong_ndim(self, xp: ModuleType):
+        m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        w2d = xp.asarray([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
+        # Non-integer correction forces the generic path where the
+        # validation lives; native backends raise for the same reason.
+        with pytest.raises((ValueError, TypeError)):
+            _ = cov(m, correction=0.5, fweights=w2d)
+        with pytest.raises((ValueError, TypeError)):
+            _ = cov(m, correction=0.5, aweights=w2d)
+
+    def test_weights_wrong_length(self, xp: ModuleType):
+        m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        w_bad = xp.asarray([1.0, 1.0])  # expected length 3
+        with pytest.raises((ValueError, RuntimeError)):
+            _ = cov(m, correction=0.5, fweights=w_bad)
+        with pytest.raises((ValueError, RuntimeError)):
+            _ = cov(m, correction=0.5, aweights=w_bad)
 
 
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False)

From 8dc3bf16f1d98b0a83e24cf4171642df13defeba Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Thu, 23 Jul 2026 14:18:19 +0200
Subject: [PATCH 14/16] Preserve torch autograd in the batched cov path

The generic `cov` implementation called `xp.asarray(m)` on its input. For
torch this detaches gradients and mutates the caller's tensor in place, so
`cov` on a batched tensor (ndim > 2, which routes to the generic path) with
`requires_grad=True` returned a detached result and silently zeroed the
input's grad.

The call is unnecessary: the delegation layer already guarantees `m` is an
array (it calls `array_namespace(m)` and reads `m.ndim`). Drop it, and add
a torch autograd regression test.
---
 src/array_api_extra/_lib/_funcs.py |  4 +++-
 tests/test_funcs.py                | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 318e1b02..3f05a171 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -270,7 +270,9 @@ def cov(
     xp: ArrayNamespace,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""
-    m = xp.asarray(m)
+    # NB: no `xp.asarray(m)` here. The delegation layer already guarantees `m`
+    # is an array (it calls `array_namespace(m)` and reads `m.ndim`), and on
+    # torch `xp.asarray` detaches gradients and mutates the caller's tensor.
     dtype = (
         xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64)
     )
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index d832e2d0..bf20c185 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -801,6 +801,22 @@ def test_weights_wrong_length(self, xp: ModuleType):
         with pytest.raises((ValueError, RuntimeError)):
             _ = cov(m, correction=0.5, aweights=w_bad)
 
+    def test_torch_autograd(self, torch: ModuleType):
+        # The batched (generic) path must not detach gradients or mutate the
+        # input tensor in place, as `xp.asarray` does on torch.
+        xp = torch
+        rng = np.random.default_rng(20260417)
+        m = xp.asarray(rng.random((4, 3, 20)), dtype=xp.float64)
+        m.requires_grad_(True)
+        # cov returns the array-api `Array` type; at runtime it is a torch
+        # tensor, so cast to access autograd attributes without type errors.
+        c = cast(Any, cov(m))  # batched -> generic path
+        assert c.requires_grad
+        assert m.requires_grad  # input tensor not mutated
+        c.sum().backward()
+        assert m.grad is not None
+        assert bool(xp.all(xp.isfinite(m.grad)))
+
 
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False)
 class TestOneHot:

From c9d3de927565cddcbf3dc6fe6c360c9b9ea23937 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Sun, 26 Jul 2026 18:37:49 +0200
Subject: [PATCH 15/16] MNT: address cov review feedback

---
 src/array_api_extra/_delegation.py |  16 ++---
 src/array_api_extra/_lib/_funcs.py |  41 +++++++++---
 tests/test_funcs.py                | 102 ++++++++++++++++++++++-------
 3 files changed, 116 insertions(+), 43 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 293c3ba5..5e8135e8 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -186,9 +186,9 @@ def cov(
     :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance
     of :math:`x_i`.
 
-    Extends ``numpy.cov`` with support for batch input and array-api
-    backends. Naming follows the array-api conventions used elsewhere in
-    this library (``axis``, ``correction``) rather than the numpy spellings
+    Extends ``numpy.cov`` with support for batch input.
+    Naming follows the array API conventions used elsewhere in
+    this library (``axis``, ``correction``) rather than the NumPy spellings
     (``rowvar``, ``bias``, ``ddof``); see Notes for the mapping.
 
     Parameters
@@ -244,11 +244,11 @@ def cov(
         numpy.cov(m, fweights=f)            -> cov(m, fweights=f)
         numpy.cov(m, aweights=a)            -> cov(m, aweights=a)
 
-    Unlike ``numpy.cov``, a ``RuntimeWarning`` for non-positive effective
-    degrees of freedom is only emitted on the unweighted path. The
-    weighted path omits the check so that lazy backends (e.g. Dask) can
-    stay lazy end-to-end; choose ``correction`` and weights such that the
-    effective normalizer is positive.
+    A ``RuntimeWarning`` is emitted for non-positive effective degrees of
+    freedom when the effective normalizer can be checked without materializing
+    a lazy array. When the normalizer itself is lazy (e.g. for weighted Dask
+    inputs), this check is skipped; choose ``correction`` and weights such that
+    it is positive.
 
     Examples
     --------
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 3f05a171..646ff45b 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -278,7 +278,8 @@ def cov(
     )
 
     m = atleast_nd(m, ndim=2, xp=xp)
-    m = xp.astype(m, dtype)
+    # Preserve the historical no-alias guarantee even when the dtype already matches.
+    m = xp.astype(m, dtype, copy=True)
 
     # Validate weight shapes (eager metadata, lazy-safe). Native backends
     # validate themselves; this covers the generic path (array-api-strict,
@@ -291,9 +292,16 @@ def cov(
         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:
+        weight_length = w_in.shape[0]
+        if (
+            weight_length is not None
+            and n_obs is not None
+            and not math.isnan(weight_length)
+            and not math.isnan(n_obs)
+            and weight_length != n_obs
+        ):
             msg = (
-                f"`{name}` has length {w_in.shape[0]} but `m` has {n_obs} observations"
+                f"`{name}` has length {weight_length} but `m` has {n_obs} observations"
             )
             raise ValueError(msg)
 
@@ -312,15 +320,9 @@ def cov(
     else:
         w = fw * aw
 
-    m_shape = eager_shape(m)
     if w is None:
         avg = xp.mean(m, axis=-1, keepdims=True)
-        fact = m_shape[-1] - correction
-        if fact <= 0:
-            warnings.warn(
-                "Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2
-            )
-            fact = 0
+        fact = eager_shape(m, axis=-1)[0] - correction
     else:
         v1 = xp.sum(w, axis=-1)
         avg = xp.sum(m * w, axis=-1, keepdims=True) / v1
@@ -329,6 +331,25 @@ def cov(
         else:
             fact = v1 - correction * xp.sum(w * aw, axis=-1) / v1
 
+    if not _compat.is_lazy_array(fact):
+        # Weights are cast to `dtype`, so a complex input produces a complex
+        # normalizer with a zero imaginary part. Complex ordering is undefined;
+        # compare its real component instead.
+        if w is not None:
+            fact_array = cast(Array, fact)
+            fact_to_check = (
+                xp.real(fact_array)
+                if xp.isdtype(fact_array.dtype, "complex floating")
+                else fact_array
+            )
+        else:
+            fact_to_check = fact
+        if fact_to_check <= 0:
+            warnings.warn(
+                "Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2
+            )
+            fact = 0
+
     m_c = m - avg
     m_w = m_c if w is None else m_c * w
     m_cT = xp.matrix_transpose(m_c)
diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index bf20c185..3e2e15cf 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -617,6 +617,27 @@ def test_complex(self, xp: ArrayNamespace):
         expect = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]], dtype=xp.complex128)
         assert_close(actual, expect)
 
+    def test_complex_with_weights(self, xp: ArrayNamespace):
+        m = np.asarray(
+            [[1 + 1j, 2 + 2j, 4 + 1j], [3 - 1j, 5 + 2j, 7 + 0j]],
+            dtype=np.complex128,
+        )
+        weights = np.asarray([1.0, 2.0, 1.0])
+        correction = 0.5  # Force the generic implementation.
+
+        weight_sum = weights.sum()
+        avg = (m * weights).sum(axis=-1, keepdims=True) / weight_sum
+        centered = m - avg
+        normalizer = weight_sum - correction * (weights**2).sum() / weight_sum
+        expected = (centered * weights) @ centered.conj().T / normalizer
+
+        actual = cov(
+            xp.asarray(m),
+            correction=correction,
+            aweights=xp.asarray(weights),
+        )
+        assert_close(actual, xp.asarray(expected))
+
     def test_empty(self, xp: ArrayNamespace):
         with warnings.catch_warnings(record=True):
             warnings.simplefilter("always", RuntimeWarning)
@@ -669,20 +690,20 @@ def test_batch(self, xp: ArrayNamespace):
         assert_close(res, xp.asarray(ref))
 
     @pytest.mark.parametrize("bias", [True, False, 0, 1])
-    def test_bias(self, xp: ModuleType, bias: bool):
+    def test_bias(self, xp: ArrayNamespace, bias: bool):
         # `bias` maps to `correction`: bias=True -> correction=0, bias=False -> 1.
         x = np.array([-2.1, -1, 4.3])
         y = np.array([3, 1.1, 0.12])
         X = np.stack((x, y), axis=0)
         ref = np.cov(X, bias=bias)
-        xp_assert_close(
+        assert_close(
             cov(xp.asarray(X, dtype=xp.float64), correction=0 if bias else 1),
             xp.asarray(ref, dtype=xp.float64),
             rtol=1e-6,
         )
 
     @pytest.mark.parametrize("bias", [True, False, 0, 1])
-    def test_bias_batch(self, xp: ModuleType, bias: bool):
+    def test_bias_batch(self, xp: ArrayNamespace, bias: bool):
         rng = np.random.default_rng(8847643423)
         batch_shape = (3, 4)
         n_var, n_obs = 3, 20
@@ -690,17 +711,17 @@ def test_bias_batch(self, xp: ModuleType, bias: bool):
         res = cov(xp.asarray(m), correction=0 if bias else 1)
         ref_list = [np.cov(m_, bias=bias) for m_ in np.reshape(m, (-1, n_var, n_obs))]
         ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
 
-    def test_correction(self, xp: ModuleType):
+    def test_correction(self, xp: ArrayNamespace):
         rng = np.random.default_rng(20260417)
         m = rng.random((3, 20))
         for correction in (0, 1, 2):
             ref = np.cov(m, ddof=correction)
             res = cov(xp.asarray(m), correction=correction)
-            xp_assert_close(res, xp.asarray(ref))
+            assert_close(res, xp.asarray(ref))
 
-    def test_correction_float(self, xp: ModuleType):
+    def test_correction_float(self, xp: ArrayNamespace):
         # Float correction: reference computed by hand (numpy.cov rejects
         # non-integer ddof; our generic path supports it).
         rng = np.random.default_rng(20260417)
@@ -709,34 +730,34 @@ def test_correction_float(self, xp: ModuleType):
         centered = m - m.mean(axis=-1, keepdims=True)
         ref = centered @ centered.T / (n - 1.5)
         res = cov(xp.asarray(m), correction=1.5)
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
 
-    def test_axis(self, xp: ModuleType):
+    def test_axis(self, xp: ArrayNamespace):
         rng = np.random.default_rng(20260417)
         m = rng.random((20, 3))  # observations on axis 0
         ref = np.cov(m, rowvar=False)
         res = cov(xp.asarray(m), axis=0)
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
         res_neg = cov(xp.asarray(m), axis=-2)
-        xp_assert_close(res_neg, xp.asarray(ref))
+        assert_close(res_neg, xp.asarray(ref))
 
-    def test_frequency_weights(self, xp: ModuleType):
+    def test_frequency_weights(self, xp: ArrayNamespace):
         rng = np.random.default_rng(20260417)
         m = rng.random((3, 10))
         fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64)
         ref = np.cov(m, fweights=fw)
         res = cov(xp.asarray(m), fweights=xp.asarray(fw))
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
 
-    def test_weights(self, xp: ModuleType):
+    def test_weights(self, xp: ArrayNamespace):
         rng = np.random.default_rng(20260417)
         m = rng.random((3, 10))
         aw = rng.random(10)
         ref = np.cov(m, aweights=aw)
         res = cov(xp.asarray(m), aweights=xp.asarray(aw))
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
 
-    def test_both_weights(self, xp: ModuleType):
+    def test_both_weights(self, xp: ArrayNamespace):
         rng = np.random.default_rng(20260417)
         m = rng.random((3, 10))
         fw = np.asarray([1, 2, 1, 3, 1, 2, 1, 1, 2, 1], dtype=np.int64)
@@ -749,9 +770,9 @@ def test_both_weights(self, xp: ModuleType):
                 fweights=xp.asarray(fw),
                 aweights=xp.asarray(aw),
             )
-            xp_assert_close(res, xp.asarray(ref))
+            assert_close(res, xp.asarray(ref))
 
-    def test_batch_with_weights(self, xp: ModuleType):
+    def test_batch_with_weights(self, xp: ArrayNamespace):
         rng = np.random.default_rng(20260417)
         batch_shape = (2, 3)
         n_var, n_obs = 3, 15
@@ -760,9 +781,9 @@ def test_batch_with_weights(self, xp: ModuleType):
         res = cov(xp.asarray(m), aweights=xp.asarray(aw))
         ref_list = [np.cov(m_, aweights=aw) for m_ in np.reshape(m, (-1, n_var, n_obs))]
         ref = np.reshape(np.stack(ref_list), (*batch_shape, n_var, n_var))
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
 
-    def test_axis_with_weights(self, xp: ModuleType):
+    def test_axis_with_weights(self, xp: ArrayNamespace):
         # axis=-2 (observations on first of 2D) combined with weights:
         # verifies that moveaxis and weight alignment cooperate.
         rng = np.random.default_rng(20260417)
@@ -776,14 +797,14 @@ def test_axis_with_weights(self, xp: ModuleType):
             fweights=xp.asarray(fw),
             aweights=xp.asarray(aw),
         )
-        xp_assert_close(res, xp.asarray(ref))
+        assert_close(res, xp.asarray(ref))
 
-    def test_axis_out_of_bounds(self, xp: ModuleType):
+    def test_axis_out_of_bounds(self, xp: ArrayNamespace):
         m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
         with pytest.raises(IndexError):
             _ = cov(m, axis=5)
 
-    def test_weights_wrong_ndim(self, xp: ModuleType):
+    def test_weights_wrong_ndim(self, xp: ArrayNamespace):
         m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
         w2d = xp.asarray([[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]])
         # Non-integer correction forces the generic path where the
@@ -793,7 +814,7 @@ def test_weights_wrong_ndim(self, xp: ModuleType):
         with pytest.raises((ValueError, TypeError)):
             _ = cov(m, correction=0.5, aweights=w2d)
 
-    def test_weights_wrong_length(self, xp: ModuleType):
+    def test_weights_wrong_length(self, xp: ArrayNamespace):
         m = xp.asarray([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
         w_bad = xp.asarray([1.0, 1.0])  # expected length 3
         with pytest.raises((ValueError, RuntimeError)):
@@ -801,18 +822,49 @@ def test_weights_wrong_length(self, xp: ModuleType):
         with pytest.raises((ValueError, RuntimeError)):
             _ = cov(m, correction=0.5, aweights=w_bad)
 
-    def test_torch_autograd(self, torch: ModuleType):
+    def test_weights_unknown_length(self, da: ArrayNamespace):
+        m_np = np.asarray([[1.0, 2.0, 3.0], [4.0, 6.0, 8.0]])
+        weights_np = np.asarray([1.0, 2.0, 3.0])
+        keep_np = np.asarray([True, False, True])
+
+        keep = da.asarray(keep_np)
+        m = da.asarray(m_np)[:, keep]
+        weights = da.asarray(weights_np)[keep]
+        assert math.isnan(m.shape[-1])
+        assert math.isnan(weights.shape[0])
+
+        actual = cov(m, aweights=weights)
+        desired = np.cov(m_np[:, keep_np], aweights=weights_np[keep_np])
+        assert_close(actual, da.asarray(desired))
+
+    def test_weights_dof_warning_eager(self):
+        xp = array_namespace(cast(Array, np.empty(0)))
+        m = xp.asarray([[1.0, 2.0], [3.0, 4.0]])
+        weights = xp.asarray([1.0, 1.0])
+
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            _ = cov(m, correction=2.5, aweights=weights)
+        assert any(
+            isinstance(warning.message, RuntimeWarning)
+            and "Degrees of freedom <= 0" in str(warning.message)
+            for warning in caught
+        )
+
+    def test_torch_autograd(self, torch: ArrayNamespace):
         # The batched (generic) path must not detach gradients or mutate the
         # input tensor in place, as `xp.asarray` does on torch.
         xp = torch
         rng = np.random.default_rng(20260417)
         m = xp.asarray(rng.random((4, 3, 20)), dtype=xp.float64)
         m.requires_grad_(True)
+        m_before = m.detach().clone()
         # cov returns the array-api `Array` type; at runtime it is a torch
         # tensor, so cast to access autograd attributes without type errors.
         c = cast(Any, cov(m))  # batched -> generic path
         assert c.requires_grad
         assert m.requires_grad  # input tensor not mutated
+        assert_equal(m.detach(), m_before)
         c.sum().backward()
         assert m.grad is not None
         assert bool(xp.all(xp.isfinite(m.grad)))

From 294d20b44c861786fbfa3bfce0bc4605740b3b22 Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Sun, 26 Jul 2026 18:43:20 +0200
Subject: [PATCH 16/16] TYP: clarify array cast in cov warning test

---
 tests/test_funcs.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_funcs.py b/tests/test_funcs.py
index 3e2e15cf..ed6f8afc 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -838,7 +838,7 @@ def test_weights_unknown_length(self, da: ArrayNamespace):
         assert_close(actual, da.asarray(desired))
 
     def test_weights_dof_warning_eager(self):
-        xp = array_namespace(cast(Array, np.empty(0)))
+        xp = array_namespace(cast(Array, cast(object, np.empty(0))))
         m = xp.asarray([[1.0, 2.0], [3.0, 4.0]])
         weights = xp.asarray([1.0, 1.0])