random_flax_module broken about numpyro

[fehiepsi]

To seed a program, you can put it under a handler

with handlers.seed(rng_seed=0):
    # random_flax_module...

[fehiepsi]

Good catch! the error means that this is run outside of numpyro context, in particular the seed handler. We should raise better errors here.

[EmanuelSommer]

Thanks for the quick reply! Could you please shortly elaborate on how to fix the issue in general i.e. how to specify the context in a generic way that does not interfere with any downstream sampling tasks? Thanks in advance :)

[EmanuelSommer]

That's what I thought, just wanted to make sure I don't overlook any subtlety :) Thanks again!