Add a warning message when using prng_key outside of seed (#1760)

* add a warning message when using prng_key outside of seed * format properly * Fix typos * make sure that the test catch warning for prng_key
pyro-ppl · Mar 23, 2024 · 9a2f618 · 9a2f618
1 parent fe7b693
commit 9a2f618
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
@@ -622,6 +622,10 @@ def prng_key():
     :return: a PRNG key of shape (2,) and dtype unit32.
     """
     if not _PYRO_STACK:
+        warnings.warn(
+            "Cannot generate JAX PRNG key outside of `seed` handler.",
+            stacklevel=find_stack_level(),
+        )
         return
 
     initial_msg = {

diff --git a/test/test_handlers.py b/test/test_handlers.py
@@ -778,7 +778,8 @@ def guide():
 
 
 def test_prng_key():
-    assert numpyro.prng_key() is None
+    with pytest.warns(Warning, match="outside of `seed`"):
+        assert numpyro.prng_key() is None
 
     with handlers.seed(rng_seed=0):
         rng_key = numpyro.prng_key()