Skip to content

gh-402: wrap NumPy RNGs to return array_api_strict Array #622

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 6 commits into from
May 30, 2025
Merged

gh-402: wrap NumPy RNGs to return array_api_strict Array #622

merged 6 commits into from
May 30, 2025

Conversation

Saransh-cpp
Copy link
Member

@Saransh-cpp Saransh-cpp commented May 14, 2025 •
edited
Loading

Description

A new Generator class to run NumPy RNG functions but wrap the final result in array_api_strict.asarray.

Refs: #402

Checks

  • Is your code passing linting?
  • Is your code passing tests?
  • Have you added additional tests (if required)?
  • Have you modified/extended the documentation (if required)?
  • Have you added a one-liner changelog entry above (if required)?

@Saransh-cpp Saransh-cpp self-assigned this May 14, 2025
@Saransh-cpp Saransh-cpp added enhancement New feature or request array-api Work is related to the Array API labels May 14, 2025
@Saransh-cpp
Copy link
Member Author

See in action - #595 (an asarray call on the output of urng is now not required).

paddyroddy
paddyroddy requested changes May 16, 2025
View reviewed changes
Copy link
Member

@paddyroddy paddyroddy left a comment

Choose a reason for hiding this comment

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

Thanks for sorting this @Saransh-cpp!

tests/test_jax.py
Comment on lines 53 to 54
assert rng.key != key
assert rvs.shape == (10_000,)
Copy link
Member

Choose a reason for hiding this comment

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

Why are these not jnp.?

glass/_array_api_utils.py

if TYPE_CHECKING:
from types import ModuleType

from jaxtyping import Array
from numpy.typing import NDArray
from array_api_strict._array_object import Array as AArray
Copy link
Member

Choose a reason for hiding this comment

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

This feels like it might get confusing

Copy link
Member

Choose a reason for hiding this comment

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

Bump

@Saransh-cpp Saransh-cpp requested a review from paddyroddy May 29, 2025 14:45
paddyroddy
paddyroddy approved these changes May 29, 2025
View reviewed changes
Copy link
Member

@paddyroddy paddyroddy left a comment

Choose a reason for hiding this comment

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

Looks good

glass/_array_api_utils.py

if TYPE_CHECKING:
from types import ModuleType

from jaxtyping import Array
from numpy.typing import NDArray
from array_api_strict._array_object import Array as AArray
Copy link
Member

Choose a reason for hiding this comment

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

Bump

@Saransh-cpp
Copy link
Member Author

@ntessore this should be ready to review whenever you are back. Merging this will un-block a bunch of PRs!

paddyroddy
paddyroddy approved these changes May 30, 2025
View reviewed changes
Copy link
Member

@paddyroddy paddyroddy left a comment

Choose a reason for hiding this comment

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

Going to merge this as needed for downstream PRs

@paddyroddy paddyroddy merged commit b0c0568 into main May 30, 2025
17 checks passed
@paddyroddy paddyroddy deleted the saransh/array_api_strict-rng branch May 30, 2025 11:21
@paddyroddy paddyroddy mentioned this pull request May 30, 2025
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@paddyroddy paddyroddy paddyroddy approved these changes

@ntessore ntessore Awaiting requested review from ntessore

Assignees

@Saransh-cpp Saransh-cpp

Labels
array-api Work is related to the Array API enhancement New feature or request
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@Saransh-cpp @paddyroddy