gh-402: add NumPy-like interface for JAX RNGs

[Saransh-cpp]

Description

This PR adds a NumPy-like interface for JAX RNGs (taken from https://github.com/glass-dev/rng-jax) and attempts to create a dispatching mechanism for the pytest fixture.

Refs: #402

Checks

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

[Saransh-cpp]

Requesting a review to start discussions/get suggestions 🙂

[paddyroddy]

@paddyroddy @ntessore this should now be ready for a review. I want to run the API design through you both before writing any tests.

The API design that I have in mind:

I think this looks nice

[ntessore]

The API design that I have in mind:

I like the design. Can we stick jax-specific functionality into a glass.jax namespace so that the optional dependency is clearly isolated?

[Saransh-cpp]

pre-commit.ci autofix

[Saransh-cpp]

The RTD failure looks unrelated.

[ntessore]

Could we test using xp and urng somewhere here? I suspect it would not look great in the logs because we use the modules itself as the labels of the xp fixture. Might make more sense to use the keys as the labels?

[Saransh-cpp]

Oops, the ids were actually not required. Everything looks nice now. See #595's test run for logs.

[Saransh-cpp]

This is now blocked by #619