From c61eea71d7edbe8110ce2a7eb8faedbccb761b7a 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/19] 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 e446d35d..2d83b8ce 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -169,7 +169,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).
 
@@ -180,16 +189,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.
 
@@ -199,6 +229,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
@@ -251,16 +298,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 2d89bf44..87546c3a 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -263,9 +263,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)
     )
@@ -273,21 +281,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 f5d756b0..bf20988e 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -675,6 +675,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 aafbb94be4edac3a953b87fdf3f6d4f7c322d9ee 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/19] 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 2d83b8ce..25076124 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -317,13 +317,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 87546c3a..567caea7 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -281,16 +281,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 16fde920f9a501ba40faa30a60aaf3aee520ac81 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/19] 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 25076124..bc9a1843 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -317,9 +317,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 567caea7..f92be343 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -286,11 +286,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 6da8bce0127c9ab22321216ea0f3608f13e9e830 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/19] 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 bf20988e..e63c9698 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -675,6 +675,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 798a21782e85e022d112cef0ec2bfb436f7d91e4 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/19] 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 f92be343..f7d33365 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -281,11 +281,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 33dac7a3056d0e47d92790c2cd1f7dc834a34dcf 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/19] 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 bc9a1843..2892dec9 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -212,11 +212,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 4609647df27dae5175b9e658eafd5f7e2aa12ea8 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/19] 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 2892dec9..957a013b 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -175,8 +175,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:
     """
@@ -214,12 +214,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.
 
@@ -237,8 +237,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
@@ -313,12 +313,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
@@ -333,15 +333,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 f7d33365..61aa51cd 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -268,8 +268,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."""
@@ -282,9 +282,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 e63c9698..d9fcb22a 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -732,7 +732,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):
@@ -740,7 +740,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):
@@ -753,8 +753,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))
 
@@ -764,7 +764,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))
@@ -780,8 +780,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 04fa2f023fbdbd1efd172821ce597b1afe346685 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/19] 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 957a013b..1ee459ae 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -309,6 +309,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 d9fcb22a..b1863511 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -790,6 +790,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 da1b5aa5cf5d58b86d7b0789ae39d7c623f95b89 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/19] 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 1ee459ae..630204c8 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -296,7 +296,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)
@@ -320,9 +320,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 61aa51cd..41a866ab 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -318,7 +318,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 00c2ca8237a0d7ff9e969f2e1ec0d15965a46387 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/19] 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 630204c8..a496b39a 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -309,20 +309,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 41a866ab..e672d7fa 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -281,6 +281,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 b1863511..e1aa66ba 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -790,18 +790,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 e115f56d768320ce42f676c3296827622f56cbd7 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/19] 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 e672d7fa..9238c96f 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -335,7 +335,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 8700e6f2e07aef31118d579b8acb5aa74ded9eae 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/19] 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 a496b39a..e518c83c 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -212,6 +212,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 41a9c6a4299bf10282bf561980aff183fee242f4 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/19] 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 e1aa66ba..4f61baf8 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -790,6 +790,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 c27a49f66cf63af343ccd167df54ab2a8b9fadda 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/19] 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 9238c96f..4e21ab1f 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -273,7 +273,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 4f61baf8..3b3bde7e 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -808,6 +808,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 17ff68df60e2204c28fca5b54632b90e757ba54a 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/19] 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 e518c83c..a4146096 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -189,9 +189,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
@@ -247,11 +247,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 4e21ab1f..3bcfea46 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -281,7 +281,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,
@@ -294,9 +295,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)
 
@@ -315,15 +323,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
@@ -332,6 +334,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 3b3bde7e..ea5c992d 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -624,6 +624,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)
@@ -676,20 +697,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
@@ -697,17 +718,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)
@@ -716,34 +737,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)
@@ -756,9 +777,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
@@ -767,9 +788,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)
@@ -783,14 +804,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
@@ -800,7 +821,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)):
@@ -808,18 +829,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 54d53157073264118b6a83564141f82dbf64c9e8 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/19] 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 ea5c992d..738f7982 100644
--- a/tests/test_funcs.py
+++ b/tests/test_funcs.py
@@ -845,7 +845,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])
 

From f76b25cd255a4581205ae01887aaf2f87edc9c64 Mon Sep 17 00:00:00 2001
From: Lucas Colley <lucas.colley8@gmail.com>
Date: Wed, 12 Aug 2026 11:28:49 +0100
Subject: [PATCH 17/19] lint

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

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index a4146096..d2668375 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -174,7 +174,7 @@ def cov(
     /,
     *,
     axis: int = -1,
-    correction: int | float = 1,
+    correction: float = 1,
     fweights: Array | None = None,
     aweights: Array | None = None,
     xp: ArrayNamespace | None = None,
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 3bcfea46..71114c90 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -267,7 +267,7 @@ def cov(
     m: Array,
     /,
     *,
-    correction: int | float = 1,
+    correction: float = 1,
     fweights: Array | None = None,
     aweights: Array | None = None,
     xp: ArrayNamespace,

From a03f59b35ddbf320fdb30eeffceaa913f6eb2f6e Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Wed, 12 Aug 2026 12:52:28 +0200
Subject: [PATCH 18/19] Apply suggestions from code review

Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>
---
 src/array_api_extra/_delegation.py | 10 +++++-----
 src/array_api_extra/_lib/_funcs.py |  5 +----
 2 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index d2668375..04649a6a 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -185,11 +185,11 @@ def cov(
     Covariance indicates the level to which two variables vary together.
     If we examine *N*-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`,
     each with *M* observations, then element :math:`C_{ij}` of the
-    :math:`N \times N` covariance matrix is the covariance of
+    :math:`N \\times N` covariance matrix is the covariance of
     :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.
+    Extends :func:`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.
@@ -202,9 +202,9 @@ def cov(
         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``
+        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``.
+        corresponds to ``rowvar=False`` in :func:`numpy.cov`.
     correction : int or float, optional
         Degrees of freedom correction: normalization divides by
         ``N - correction`` (for unweighted input). Default: ``1``, which
@@ -233,7 +233,7 @@ def cov(
     Returns
     -------
     array
-        An array having shape (..., N, N) whose innermost two dimensions represent
+        An array having shape ``(..., N, N)`` whose innermost two dimensions represent
         the covariance matrix of the variables.
 
     Notes
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index 71114c90..b0fae800 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -284,10 +284,7 @@ def cov(
     # 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,
-    # sparse, and the dask+weights fallback where the native check is
-    # bypassed to preserve laziness).
+    # Validate weight shapes (eager metadata, lazy-safe).
     n_obs = m.shape[-1]
     for name, w_in in (("fweights", fweights), ("aweights", aweights)):
         if w_in is None:

From 755f39c8b6447a610bd86ab2f6a619da23e0aa4e Mon Sep 17 00:00:00 2001
From: Bru <b.aristimunha@gmail.com>
Date: Wed, 12 Aug 2026 13:10:38 +0200
Subject: [PATCH 19/19] MNT: address cov review comments

- link numpy/torch functions with intersphinx in the cov docstring
  (torch added to the intersphinx mapping)
- add canonical examples for correction, fweights and aweights
- explain the int(correction) cast: torch.cov rejects integer-valued
  floats such as 1.0 at runtime, so typing.cast is not enough
- simplify the integer-correction check to float(correction).is_integer()
- comment that the NaN shape checks account for Dask reporting unknown
  dimensions as NaN instead of None
---
 docs/conf.py                       |  1 +
 src/array_api_extra/_delegation.py | 50 +++++++++++++++++++++++-------
 src/array_api_extra/_lib/_funcs.py |  2 ++
 3 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/docs/conf.py b/docs/conf.py
index d420b5aa..1bda1170 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -60,6 +60,7 @@
     "numpy": ("https://numpy.org/doc/stable", None),
     "jax": ("https://docs.jax.dev/en/latest", None),
     "pytest": ("https://docs.pytest.org/en/stable/", None),
+    "torch": ("https://docs.pytorch.org/docs/stable", None),
 }
 
 nitpick_ignore = [
diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py
index 04649a6a..c91517ca 100644
--- a/src/array_api_extra/_delegation.py
+++ b/src/array_api_extra/_delegation.py
@@ -208,25 +208,26 @@ def 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
+        gives the unbiased estimate (matches :func:`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``.
+        normalization). Corresponds to ``ddof`` in :func:`numpy.cov` and to
+        ``correction`` in :func:`numpy.var`/:func:`numpy.std` and
+        :func:`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.
+        generic implementation because :func:`numpy.cov`'s ``ddof`` and
+        :func:`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
-        ``numpy.cov``/``torch.cov``.
+        :func:`numpy.cov`/:func:`torch.cov`.
     aweights : array, optional
         1-D array of observation-vector weights (analytic weights). Larger
         values mark more important observations. Same as ``aweights`` in
-        ``numpy.cov``/``torch.cov``.
+        :func:`numpy.cov`/:func:`torch.cov`.
     xp : array_namespace, optional
         The standard-compatible namespace for `m`. Default: infer.
 
@@ -238,7 +239,7 @@ def cov(
 
     Notes
     -----
-    Mapping from ``numpy.cov`` to this function::
+    Mapping from :func:`numpy.cov` to this function::
 
         numpy.cov(m, rowvar=True)           -> cov(m, axis=-1)   # default
         numpy.cov(m, rowvar=False)          -> cov(m, axis=-2)
@@ -300,6 +301,30 @@ def cov(
             [ -4.286     ,   2.14413333]],
            [[ 46.84      , -17.144     ],
             [-17.144     ,   8.57653333]]], dtype=array_api_strict.float64)
+
+    The normalization can be adjusted with `correction`, and observations
+    can be weighted with integer frequencies `fweights` or importance
+    weights `aweights`:
+
+    >>> x = xp.asarray([0., 1., 2., 3., 4.])
+    >>> xpx.cov(x, xp=xp)  # unbiased variance: divide by N - 1
+    Array(2.5, dtype=array_api_strict.float64)
+    >>> xpx.cov(x, correction=0, xp=xp)  # biased variance: divide by N
+    Array(2., dtype=array_api_strict.float64)
+
+    Giving the two extreme observations frequency 2 via `fweights` is
+    equivalent to repeating them in `x`:
+
+    >>> xpx.cov(x, fweights=xp.asarray([2, 1, 1, 1, 2]), xp=xp)
+    Array(3., dtype=array_api_strict.float64)
+    >>> xpx.cov(xp.asarray([0., 0., 1., 2., 3., 4., 4.]), xp=xp)
+    Array(3., dtype=array_api_strict.float64)
+
+    `aweights` instead adjusts the relative importance of observations,
+    here down-weighting the two extremes:
+
+    >>> xpx.cov(x, aweights=xp.asarray([0.5, 1., 1., 1., 0.5]), xp=xp)
+    Array(1.92, dtype=array_api_strict.float64)
     """
 
     if xp is None:
@@ -319,14 +344,17 @@ def cov(
     # `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()
+    integer_correction = float(correction).is_integer()
     has_weights = fweights is not None or aweights is not None
 
     if m.ndim <= 2 and integer_correction:
+        # Not just for static typing: `correction` may be an integer-valued
+        # float such as 1.0, which `torch.cov` rejects at runtime.
+        int_correction = int(correction)
         if is_torch_namespace(xp):
             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)
+            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
@@ -339,7 +367,7 @@ def cov(
         ):
             return xp.cov(
                 m,
-                ddof=int(correction),
+                ddof=int_correction,
                 fweights=fweights,
                 aweights=aweights,
             )
diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py
index b0fae800..4629461e 100644
--- a/src/array_api_extra/_lib/_funcs.py
+++ b/src/array_api_extra/_lib/_funcs.py
@@ -293,6 +293,8 @@ def cov(
             msg = f"`{name}` must be 1-D, got ndim={w_in.ndim}"
             raise ValueError(msg)
         weight_length = w_in.shape[0]
+        # Unknown dims are `None` per the standard; Dask non-standardly
+        # reports them as NaN, hence the `isnan` checks below.
         if (
             weight_length is not None
             and n_obs is not None