Skip to content

Commit

Permalink
add a note to seed handler (#1759)
Browse files Browse the repository at this point in the history
  • Loading branch information
fehiepsi committed Mar 23, 2024
1 parent 9a2f618 commit 5cad57d
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions numpyro/handlers.py
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,14 @@ class seed(Messenger):
users need to use `seed` as a contextmanager to generate samples from
distributions or as a decorator for their model callable (See below).
.. note::
The seed handler has a mutable attribute `rng_key` which keeps changing after
each `sample` call. Hence an instance of this class (e.g. `seed(model, rng_seed=0)`)
might create tracer-leaks when jitted. A solution is to close the instance in a
function, e.g., `seeded_model = lambda *args: seed(model, rng_seed=0)(*args)`.
This `seeded_model` can be jitted.
**Example:**
.. doctest::
Expand Down

0 comments on commit 5cad57d

Please sign in to comment.