ENH: cov: expose correction and weights parameters by bruAristimunha · Pull Request #690 · data-apis/array-api-extra

bruAristimunha · 2026-04-17T10:43:02Z

Resolves #688.

Summary

Adds axis, correction, frequency_weights, and weights parameters to xpx.cov, unlocking the degrees-of-freedom and weighted variants that numpy.cov and torch.cov already support.
Naming follows array-api conventions (axis, correction) used elsewhere in this library rather than numpy's (rowvar, bias, ddof). The docstring includes a one-to-one mapping for users migrating from numpy.cov.

Design

The delegation moves observations to the last axis via xp.moveaxis, which collapses rowvar out of backend dispatch entirely — only ddof (numpy/cupy/dask/jax) vs correction (torch) differs between branches.

Fallbacks to the generic implementation (_funcs.cov):

m.ndim > 2 (batched input, not supported by any native).
Non-integer correction (rejected by numpy.cov's ddof).
Dask with weights — dask.array.cov forces .compute() on a lazy 0-D scalar via its internal if fact <= 0 check. The generic path stays fully lazy because its weighted branch doesn't compare fact to zero (noted in docstring).

Weighted formula in _funcs.cov matches numpy's (algebraically): c = (m_c · w) @ m_c.T / (v1 - correction · v2 / v1).

Tests

New TestCov cases validate against np.cov as reference:

test_correction (integer ddof)
test_correction_float (generic-path-only, hand-computed reference)
test_axis / test_axis_with_weights / test_axis_out_of_bounds
test_frequency_weights / test_weights / test_both_weights
test_batch_with_weights

Test plan

pytest tests/test_funcs.py::TestCov — 126 passed across numpy, torch, jax, dask, array-api-strict
pytest tests/test_funcs.py full — 4263 passed, 0 failed
lefthook run pre-commit — ruff, numpydoc, mypy, pyright, typos all green
Dask laziness verified — lazy_xp_function(cov) asserts 0 .compute() calls, holds for weighted path via the fallback

betatim · 2026-04-20T09:16:47Z

It looks like the cov you are adding follows the pytorch signature, can you explain a bit why you chose that? In my PR I thought following the Numpy API makes sense because it seems that most libraries use that.

The PR description mentions that other functions in this library already use correction and axis. Which is a good reason to also do it here? Interested in your thinking.

bruAristimunha · 2026-04-20T09:45:06Z

Hey @betatim!

This was a little hard decision that I had to make, but I can be more strict with numpy if you prefer.

I basically looked at what was already implemented on the API array and how they handle the parameter names that I was trying to implement.

Like, for each parameter that I was trying to introduce, I checked how it was made in the past here from numpy to: the bias, the rowvar, the ddof, the fweights, and the aweights.

Basically, for the bias, ddof to become correction, I notice that in the functions xp.var, xp.std, and think xp.sum, they change the default names to the array api specification name.

https://data-apis.org/array-api/latest/API_specification/generated/array_api.var.html
https://data-apis.org/array-api/latest/API_specification/generated/array_api.std.html

There was a discussion on how to use correction instead of bias+ddof on these functions. Here was introduced data-apis/array-api#10, and then, later, they made some interesting discussions here: data-apis/array-api#695; it was @kgryte who led the discussion.

For the case of the rowvar becoming the axis, I just follow the signature of the other functions. seems like the axis was how they followed.

And for the frequency_weights and weights, it was my experience in Pyriemann that made the decisions. I think the only place that I remember using something similar was the statsmodels (freq_weights, var_weights) that uses https://www.statsmodels.org/stable/generated/statsmodels.genmod.generalized_linear_model.GLM.html#statsmodels.genmod.generalized_linear_model.GLM.freq_weights

I think in scikit you guys use sample_weight more, but I can accommodate any request about this.

betatim

What is your thinking on validating the weights passed in? Things like checking the shapes make sense, that they are all positive (is this actually required? how does it fit with being lazy?)

bruAristimunha · 2026-04-20T09:55:43Z

I liked this idea a lot @betatim! I think it will make the check in the library that use api array extra much lighter.

bruAristimunha · 2026-04-20T10:12:36Z

FYI @qbarthelemy and @agramfort

betatim · 2026-04-20T11:59:18Z

Thanks a lot for the detailed answer in #690 (comment) - I didn't realise there was precedent for using correction in functions like var. I think it makes sense to copy that and use correction for cov as well. Worth making the translation!

What is the "temporary deployed" thing that keeps happening?

bruAristimunha · 2026-04-20T12:01:01Z

it is not me @betatim, i think it something that @lucascolley is pushing in pushing here: #699

bruAristimunha · 2026-04-20T12:01:58Z

Happy that you liked the response @betatim :)

I think I addressed all the points from you and @qbarthelemy, can we merge?

lucascolley · 2026-04-20T12:07:24Z

What is the "temporary deployed" thing that keeps happening?

fixed in bd3652a

lucascolley

I took an initial look, seems pretty good!

One high level comment @bruAristimunha — could you demonstrate that this works as expected when used in a branch of sklearn? You should be able to change https://github.com/scikit-learn/scikit-learn/blob/06aded051fe6c7c9970b7e13c3669f952a799831/maint_tools/vendor_array_api_extra.sh#L8-L9 to point to this branch and commit hash.

bruAristimunha · 2026-04-20T14:11:36Z

hey @betatim,

As you have the first covariance PR on scikit, can you help with this small test as requested by @lucascolley?

One high level comment @bruAristimunha — could you demonstrate that this works as expected when used in a branch of sklearn? You should be able to change https://github.com/scikit-learn/scikit-learn/blob/06aded051fe6c7c9970b7e13c3669f952a799831/maint_tools/vendor_array_api_extra.sh#L8-L9 to point to this branch and commit hash.

bruAristimunha · 2026-04-20T16:27:29Z

hey @lucascolley,

I made in my branch that was built on top of @betatim's work for scikit-learn first covariance, you can check more here: scikit-learn/scikit-learn#33600

lucascolley · 2026-04-20T17:34:58Z

hey @lucascolley,

I made in my branch that was built on top of @betatim's work for scikit-learn first covariance, you can check more here: scikit-learn/scikit-learn#33600

thanks! Would be great if you could take a look, Tim

Resolves data-apis#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.

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>

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.

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%.

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.

Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>

- 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

lucascolley

thanks @bruAristimunha for the perseverance, I think this is ready to merge now — nice work!

Thanks also @qbarthelemy @betatim for the reviews!

BTW I added pyRiemann to the list of users in the README, hope you don't mind 😉 feel free to shout if you hit any problems, or have any other functions that you would like to see added.

lucascolley · 2026-08-12T17:44:32Z

@all-contributors please add @bruAristimunha for examples

allcontributors · 2026-08-12T17:44:54Z

@lucascolley

I've put up a pull request to add @bruAristimunha! 🎉

lucascolley · 2026-08-12T17:46:17Z

@all-contributors please add @betatim for review

allcontributors · 2026-08-12T17:46:39Z

@lucascolley

I've put up a pull request to add @betatim! 🎉

bruAristimunha mentioned this pull request Apr 17, 2026

Add bias keyword argument to cov #691

Closed

betatim reviewed Apr 20, 2026

View reviewed changes

Comment thread src/array_api_extra/_delegation.py Outdated

betatim reviewed Apr 20, 2026

View reviewed changes

Comment thread src/array_api_extra/_delegation.py Outdated

betatim reviewed Apr 20, 2026

View reviewed changes

Comment thread src/array_api_extra/_lib/_funcs.py Outdated

betatim reviewed Apr 20, 2026

View reviewed changes

bruAristimunha force-pushed the cov_parameters branch from 83b7e1b to d9701e0 Compare April 20, 2026 10:10

bruAristimunha temporarily deployed to ci-checks April 20, 2026 10:26 — with GitHub Actions Inactive

bruAristimunha deployed to ci-checks April 20, 2026 10:26 — with GitHub Actions Active

bruAristimunha temporarily deployed to ci-checks April 20, 2026 10:26 — with GitHub Actions Inactive

lucascolley added the enhancement New feature or request label Apr 20, 2026

qbarthelemy reviewed Apr 20, 2026

View reviewed changes

Comment thread src/array_api_extra/_lib/_funcs.py Outdated

qbarthelemy reviewed Apr 20, 2026

View reviewed changes

Comment thread src/array_api_extra/_delegation.py Outdated

lucascolley changed the title ~~ENH: expose correction and weights parameters in cov~~ ENH: cov: expose correction and weights parameters Apr 20, 2026

lucascolley reviewed Apr 20, 2026

View reviewed changes

qbarthelemy reviewed Apr 20, 2026

View reviewed changes

Comment thread src/array_api_extra/_lib/_funcs.py Outdated

lucascolley added this to the 0.10.2 milestone Apr 24, 2026

lucascolley added this to the 0.11.2 milestone Aug 12, 2026

lucascolley self-requested a review August 12, 2026 11:32

bruAristimunha and others added 19 commits August 12, 2026 17:47

MNT: drop device= in cov weights

aafbb94

STY: formatter

16fde92

TST: add bias tests from data-apis#691

6da8bce

Update _funcs.py

798a217

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>

Update _delegation.py

33dac7a

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>

MNT: rename weights params to fweights/aweights

4609647

ENH: validate weights shape in cov

04fa2f0

MNT: address lucascolley review

da1b5aa

MNT: move weights validation to generic cov

00c2ca8

Update _funcs.py

e115f56

Co-authored-by: Quentin Barthélemy <q.barthelemy@gmail.com>

DOC: explain non-integer correction use cases in cov

8700e6f

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.

TST: cover weight validation error paths in cov

41a9c6a

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%.

MNT: address cov review feedback

17ff68d

TYP: clarify array cast in cov warning test

54d5315

lint

f76b25c

Apply suggestions from code review

a03f59b

Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>

lucascolley force-pushed the cov_parameters branch from a34ba84 to 755f39c Compare August 12, 2026 16:47

lucascolley approved these changes Aug 12, 2026

View reviewed changes

lucascolley merged commit f903723 into data-apis:main Aug 12, 2026
60 of 61 checks passed

allcontributors Bot mentioned this pull request Aug 12, 2026

docs: add bruAristimunha as a contributor for example #912

Merged

allcontributors Bot mentioned this pull request Aug 12, 2026

docs: add betatim as a contributor for review #914

Merged

Uh oh!

Conversation

bruAristimunha commented Apr 17, 2026

Summary

Design

Tests

Test plan

Uh oh!

betatim commented Apr 20, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

betatim left a comment

Choose a reason for hiding this comment

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

Uh oh!

Uh oh!

betatim commented Apr 20, 2026

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

lucascolley commented Apr 20, 2026

Uh oh!

lucascolley left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

bruAristimunha commented Apr 20, 2026

Uh oh!

lucascolley commented Apr 20, 2026

Uh oh!

Uh oh!

lucascolley left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lucascolley commented Aug 12, 2026

Uh oh!

allcontributors Bot commented Aug 12, 2026

Uh oh!

lucascolley commented Aug 12, 2026

Uh oh!

allcontributors Bot commented Aug 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants