Allow passing key as seed

liesel-devs / liesel

A probabilistic programming framework

https://liesel-project.org

MIT License

38 stars 2 forks source link

Allow passing key as seed #156

Closed Seb-Lorek closed 8 months ago

Seb-Lorek commented 8 months ago

Simple proposal to allow seed argument to be a key generated from jax.random.PRNGKey in the initialization of class EngineBuilder. Unfortunately type testing for the key is not really strict with jax.Array. However using jnp.issubdtype(seed.dtype, jnp.uint32) makes it more complicated. Stricter type testing seems only possible with jax.random.key currently, see #16716 and #16781 for current development on this.

jobrachem commented 8 months ago

I just merged the PR. Thanks @Seb-Lorek and congrats on your first code contribution to the liesel repo! 🎉 😊