Skip to content

Conversation

@amarjeetsingh146
Copy link

Summary

This PR updates the pseudorandom number documentation to clarify the
comparison between NumPy and JAX PRNG designs.

Details

  • Notes that modern NumPy encourages use of local generators via
    numpy.random.default_rng
  • Clarifies that numpy.random.seed() refers to legacy global-state usage
  • Preserves the explanation of JAX’s explicit key-based PRNG design

Motivation

The previous wording implied that NumPy PRNG relies solely on global state,
which is no longer the recommended usage in modern NumPy.

Fixes #24927

@google-cla
Copy link

google-cla bot commented Jan 3, 2026

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @amarjeetsingh146, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request improves the clarity of the random numbers guide by updating the comparison between NumPy and JAX's pseudo-random number generation (PRNG) designs. It addresses the outdated implication that NumPy solely relies on global state, highlighting modern NumPy's recommended use of local generators while maintaining the detailed explanation of JAX's explicit key-based approach. This ensures the documentation accurately reflects current best practices and provides a more precise understanding of both frameworks.

Highlights

  • Clarified NumPy PRNG Usage: The documentation now distinguishes between modern NumPy's recommended use of local generator objects via numpy.random.default_rng and the legacy global-state usage of numpy.random.seed().
  • Enhanced JAX PRNG Comparison: The explanation of JAX's explicit key-based PRNG design has been refined to highlight its benefits for reproducibility, parallelism, and vectorization, allowing for a clearer comparison with modern NumPy generators.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request updates the documentation on pseudorandom numbers to clarify the comparison between NumPy and JAX PRNG designs, particularly highlighting modern NumPy's use of local generators. The changes are a good improvement, but I've found a few typos and inconsistencies in the new text that should be corrected for clarity and consistency.

@jakevdp
Copy link
Collaborator

jakevdp commented Jan 3, 2026

Thanks for looking at this – I don't think the updated text does enough to emphasize the distinction between NumPy's model (where the PRNG state is updated implicitly, whether that state is local or global) and JAX's model (where the state must be updated explicitly).

@amarjeetsingh146
Copy link
Author

Thanks for looking at this – I don't think the updated text does enough to emphasize the distinction between NumPy's model (where the PRNG state is updated implicitly, whether that state is local or global) and JAX's model (where the state must be updated explicitly).

Thanks for the clarification — that’s a helpful distinction to call out more explicitly.

I agree that the key difference isn’t just global vs local state, but implicit state mutation in NumPy-style APIs versus explicit state threading in JAX’s PRNG model. I’ll revise the text to emphasize that NumPy updates PRNG state implicitly (regardless of scope), while JAX requires users to explicitly pass and update keys, making randomness a pure, explicit part of the computation.

I’ll push an update shortly that highlights this contrast more directly.

@amarjeetsingh146
Copy link
Author

I’ve pushed an update that consistently emphasizes implicit vs explicit PRNG state mutation throughout the section, including a small wording tweak in the multi-device discussion.
Please let me know if this addresses the concern.

@jakevdp
Copy link
Collaborator

jakevdp commented Jan 5, 2026

The new text does not match the older unmodified code snippets, which makes it hard to follow.


```{code-cell}
import numpy as np
np.random.seed(0)
Copy link
Collaborator

Choose a reason for hiding this comment

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

It's strange that we show an example of np.random.seed right after saying it's a legacy function that shouldn't be used.

@jakevdp jakevdp self-assigned this Jan 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Rework comparisons in the docs of JAX vs. NumPy PRNG

2 participants