gh-1145: Default to wrapping rng if backend is not known - #1169
Conversation
7b789f3 to
e430a79
Compare
connoraird
changed the title
| import glass.jax # noqa: PLC0415 | ||
|
|
||
| return glass.jax.Generator(seed=seed) | ||
| return glass.jax.Generator(seed=seed) # ty: ignore[invalid-return-type] |
There was a problem hiding this comment.
Should we be changing the output type above from UnifiedGenerator?
There was a problem hiding this comment.
Change it to what? An alternative is to keep UnifiedGenerator as it was and give this new protocol class a different name?
There was a problem hiding this comment.
I've updated the name of the protocol
There was a problem hiding this comment.
Will have a look, but yes the alternative was some long-winded UnifiedGenerator | glass.jax.Generator | ... etc.
|
@paddyroddy I've made some significant changes after discovering a bug when adding the new tests with JAX. |
Okay, will have a look |
| assert isinstance(rng, glass.rng.Generator) | ||
|
|
||
|
|
||
| @pytest.mark.skipif(not HAVE_ARRAY_API_STRICT, reason="test requires array_api_strict") |
| """Draw samples from a multinomial distribution.""" | ||
| # Ensure arrays are jax arrays | ||
| n = jnp.asarray(n) | ||
| pvals = jnp.asarray(pvals) |
There was a problem hiding this comment.
Interesting that it's only these two methods that needed it
There was a problem hiding this comment.
In this case, it's needed because of the way we index pvals below pvals[..., 0]. I suppose n may not be necessary but it's a bit more complete.
There was a problem hiding this comment.
Although, maybe I'm misremembering as both instances are passing into _shape so that might be it.
There was a problem hiding this comment.
Remind me why this change is needed now? Is it because we are now actually testing with JAX?
|
I've removed the Protocol class as I don't think it was actually doing anything. It definitely wasn't doing what I wanted which was restricting what RNGs could be passed to |
Okay makes sense. Would be good if @ntessore could explain them to us a bit more 😅 |
|
Could have been a fluke, but the examples took more than 20 mins and failed. Looks like they usually take about 15 mins. Have re-run, otherwise we might have had a regression. |
paddyroddy
left a comment
paddyroddy
left a comment
There was a problem hiding this comment.
Looks good, see one comment on examples timing
| if rng is None: | ||
| if xp.__name__ == "jax.numpy": | ||
| import glass.jax # noqa: PLC0415 | ||
|
|
||
| self.rng = glass.jax.Generator(seed=seed) | ||
| else: | ||
| import numpy as np # noqa: PLC0415 | ||
|
|
||
| self.rng = np.random.default_rng(seed=seed) | ||
| else: | ||
| self.rng = rng |
There was a problem hiding this comment.
I think this logic is nice and clear now
|
Examples took 21:00 |
2 participants
Description
glass.rng.Generatorto allow it to wrap more than just NumPy RNGs.glass.rng.Generatordefines the required methods.Closes: #1145
Changelog entry
Changed:
glass.rng.Generatorclass now wraps returned values ensuring all methods return the correct array backend.Checks