ENH: cov: expose correction and weights parameters

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,
+    fweights: Array | None = None,
+    aweights: Array | None = None,
+    xp: ArrayNamespace | None = None,
+) -> Array:
     """
     Estimate a covariance matrix (or a stack of covariance matrices).
 
@@ -177,16 +186,44 @@ 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.
+    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``.
+        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
+        ``numpy.cov``/``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``.
     xp : array_namespace, optional
         The standard-compatible namespace for `m`. Default: infer.
 
@@ -196,6 +233,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, fweights=f)
+        numpy.cov(m, aweights=a)            -> cov(m, aweights=a)
+
+    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
     --------
     >>> import array_api_strict as xp
@@ -247,18 +301,54 @@ 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)
+        xp = array_namespace(m, fweights, aweights)
+
+    # 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 = fweights is not None or aweights is not None
+
+    if m.ndim <= 2 and integer_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)
+        # `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=fweights,
+                aweights=aweights,
+            )
+
+    return _funcs.cov(
+        m,
+        correction=correction,
+        fweights=fweights,
+        aweights=aweights,
+        xp=xp,
+    )
 
 
 def create_diagonal(

src/array_api_extra/_lib/_funcs.py

@@ -260,31 +260,102 @@ 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,
+    fweights: Array | None = None,
+    aweights: Array | None = None,
+    xp: ArrayNamespace,
+) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""
-    m = xp.asarray(m, copy=True)
+    # 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)
     )
 
     m = atleast_nd(m, ndim=2, xp=xp)
-    m = xp.astype(m, dtype)
-
-    avg = xp.mean(m, axis=-1, keepdims=True)
-
-    m_shape = eager_shape(m)
-    fact = m_shape[-1] - 1
+    # 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).
+    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)
+        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 {weight_length} but `m` has {n_obs} observations"
+            )
+            raise ValueError(msg)
 
-    if fact <= 0:
-        warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2)
-        fact = 0
+    fw = None
+    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:
+        w = aw
+    elif aw is None:
+        w = fw
+    else:
+        w = fw * aw
 
-    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 = eager_shape(m, axis=-1)[0] - correction
+    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
+
+    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)
+    if xp.isdtype(m_cT.dtype, "complex floating"):
+        m_cT = xp.conj(m_cT)
+    c = m_w @ m_cT / fact
     axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
     return xp.squeeze(c, axis=axes)