feat: convenience construction of a threefry

[nstarman]

No description provided.

[nstarman]

I think the test failures are unrelated.

[patrick-kidger]

Hmm, I think the test failures are probably due to a change in JAX. If you are able to fix that then I'd be happy to merge this.

[nstarman]

I can't reproduce the error locally. Something in zeros.

[patrick-kidger]

Which version of JAX are you using? I suspect it's probably due to an update on their end.
(I haven't checked this myself yet.)

[nstarman]

Which version of JAX are you using? I suspect it's probably due to an update on their end. (I haven't checked this myself yet.)

macOS
Python 3.12
jax[cpu] 0.4.25. I tried bumping to most recent jax in a clean environment.

[patrick-kidger]

Okay, so possibly an issue with the latest version of JAX! If this passes with an earlier version then that will be the culprit, and we'll need to figure out how to use JAX's changes.

[nstarman]

Now that jax has a key type that isn't just an Array and can be used for annotations, this PR isn't necessary.