support passing jax.random.PRNGKey inputs in jax
#651
Conversation
| if isinstance(seed, jax.Array): | ||
| return seed | ||
| else: | ||
| return draw_seed(seed) |
There was a problem hiding this comment.
Shouldn't we always convert the seed to a jax PRNGKey in case the backend is JAX? I am just thinking about the advantage of the passing any other seed type
There was a problem hiding this comment.
There are no particular performance advantages (or any other advantages) to passing a PRNGKey over a JAX array. The PRNGKey class is syntactic sugar.
There was a problem hiding this comment.
Sorry for the confusion. I was referring to the other types (else branch in this case), like a single integer
There was a problem hiding this comment.
oh yeah the single integers are probably getting broadcast from x -> [x, x], which isn't obviously any better or worse than x -> [0, x] which is what PRNGKey seems to do in practice, though it might do something different with non-default rng implementations?
|
added a unit test and made sure it passed |
result will already be an rng
No description provided.