Skip to content

Add correction argument to jnp.nanvar and jnp.nanstd - #41696

Merged
copybara-service[bot] merged 1 commit into
jax-ml:mainfrom
jayzuccarelli:nanvar-correction
Oct 19, 2026
Merged

copybara-service[bot] merged 1 commit into
jax-ml:mainfrom
jayzuccarelli:nanvar-correction

Conversation

@jayzuccarelli

Copy link
Copy Markdown
Contributor

NumPy 2.0 added correction to nanvar/nanstd as an alias for ddof, same as var/std. This mirrors the existing jnp.var/jnp.std handling (#21088): the public functions resolve correction/ddof and raise if both are given, and the jitted work moves into _nanvar/_nanstd so an explicit ddof=0 isn't traced.

Also drops nanvar/nanstd from the known-missing list in testWrappedSignaturesMatch, and adds the mean parameter the stubs were missing.

Comment thread jax/_src/numpy/reductions.py Resolved
Comment thread jax/_src/numpy/reductions.py Resolved

@jakevdp jakevdp left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks pretty good! Just one more minor comment

Comment thread tests/lax_numpy_reducers_test.py Resolved
@jakevdp jakevdp self-assigned this Oct 16, 2026
@jakevdp

jakevdp commented Oct 19, 2026

Copy link
Copy Markdown
Collaborator

I think the failures here are unrelated. Could you please squash your changes into a single commit and rebase on the updated main branch? Thanks!

NumPy 2.0 added `correction` to `nanvar`/`nanstd` as an alias for
`ddof`, same as `var`/`std`. This mirrors the existing `jnp.var`/`jnp.std`
handling (jax-ml#21088): the public functions resolve `correction`/`ddof` and
raise if both are given, and the jitted work moves into `_nanvar`/`_nanstd`
so an explicit `ddof=0` isn't traced.

Also drops `nanvar`/`nanstd` from the known-missing list in
`testWrappedSignaturesMatch`, and adds the `mean` parameter the stubs were
missing.

@jakevdp jakevdp left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@google-ml-butler google-ml-butler Bot added the pull ready Ready for copybara import and testing label Oct 19, 2026
@copybara-service
copybara-service Bot merged commit 8d98d63 into jax-ml:main Oct 19, 2026
58 of 59 checks passed
@jayzuccarelli
jayzuccarelli deleted the nanvar-correction branch October 19, 2026 18:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pull ready Ready for copybara import and testing

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants