We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Document behavior of primitive statements in README
Hello! I'll preface this issue by stating that I'm new to numpyro, so there's a significant chance that I'm not using `numpyro.sample` correctly.

When building models, I want to test draws from my random variables to make sure I'm passing parameters correctly, getting the kwargs correct, etc.; I'll lift a line out of a model function, [like this one](https://github.com/pyro-ppl/numpyro/blob/master/examples/baseball.py#L111), and run it in an IPython window to see if I get an error. It seems like I'm unable to do this by itself in an IPython console.

## Minimal example

```
In [1]: import numpyro

In [2]: import numpyro.distributions as d

In [3]: numpyro.__version__
Out[3]: '0.2.0'

In [4]: numpyro.sample("x", d.Normal(0, 1))
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-4-21ce96e72ad0> in <module>
----> 1 numpyro.sample("x", d.Normal(0, 1))

~/miniconda3/envs/fusion/lib/python3.7/site-packages/numpyro/primitives.py in sample(name, fn, obs, sample_shape)
     45     # if there are no active Messengers, we just draw a sample and return it as expected:
     46     if not _PYRO_STACK:
---> 47         return fn(sample_shape=sample_shape)
     48
     49     # Otherwise, we initialize a message...

~/miniconda3/envs/fusion/lib/python3.7/site-packages/numpyro/distributions/distribution.py in __call__(self, *args, **kwargs)
    161
    162     def __call__(self, *args, **kwargs):
--> 163         key = kwargs.pop('random_state')
    164         sample_intermediates = kwargs.pop('sample_intermediates', False)
    165         if sample_intermediates:

KeyError: 'random_state'
```

## Thoughts
I can see that in [numpyro.sample](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/primitives.py#L47), we return `fn(sample_shape=sample_shape)`, which in this case attempts to call `dist.Normal(0, 1)(sample_shape=sample_shape)`. However, looking at [distributions.py](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/distribution.py#L163), it seems that `Distribution.sample` expects a `random_state` kwarg that isn't getting passed.

When I do the following, everything is fine:

```
In [5]: import jax.random as random

In [6]: d.Normal(0, 1).sample(random.PRNGKey(12))
/Users/chtu8001/miniconda3/envs/fusion/lib/python3.7/site-packages/jax/lib/xla_bridge.py:114: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Out[6]: DeviceArray(-0.5135873, dtype=float32)

In [7]: d.Normal(0, 1).sample(random.PRNGKey(12), sample_shape=(4, ))
Out[7]:
DeviceArray([-1.31179953, -0.70821768,  0.18907626, -1.09385514],
            dtype=float32)
```

so I expect that we can't use `numpyro.sample` outside a model (although it's not totally clear to me how that is defined, something to do with Messengers and PYRO_STACK). I'm wondering if this is by design and I should just use the second, working example, or if I'm misunderstanding how `numpyro.sample` should be used.

Thanks!

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `numpyro/primitives.py`
Content:
```
1 from collections import namedtuple
2 import functools
3 
4 import jax
5 from jax import lax
6 
7 import numpyro
8 from numpyro.distributions.discrete import PRNGIdentity
9 
10 _PYRO_STACK = []
11 
12 
13 CondIndepStackFrame = namedtuple('CondIndepStackFrame', ['name', 'dim', 'size'])
14 
15 
16 def apply_stack(msg):
17     pointer = 0
18     for pointer, handler in enumerate(reversed(_PYRO_STACK)):
19         handler.process_message(msg)
20         # When a Messenger sets the "stop" field of a message,
21         # it prevents any Messengers above it on the stack from being applied.
22         if msg.get("stop"):
23             break
24     if msg['value'] is None:
25         if msg['type'] == 'sample':
26             msg['value'], msg['intermediates'] = msg['fn'](*msg['args'],
27                                                            sample_intermediates=True,
28                                                            **msg['kwargs'])
29         else:
30             msg['value'] = msg['fn'](*msg['args'], **msg['kwargs'])
31 
32     # A Messenger that sets msg["stop"] == True also prevents application
33     # of postprocess_message by Messengers above it on the stack
34     # via the pointer variable from the process_message loop
35     for handler in _PYRO_STACK[-pointer-1:]:
36         handler.postprocess_message(msg)
37     return msg
38 
39 
40 class Messenger(object):
41     def __init__(self, fn=None):
42         self.fn = fn
43         functools.update_wrapper(self, fn, updated=[])
44 
45     def __enter__(self):
46         _PYRO_STACK.append(self)
47 
48     def __exit__(self, *args, **kwargs):
49         assert _PYRO_STACK[-1] is self
50         _PYRO_STACK.pop()
51 
52     def process_message(self, msg):
53         pass
54 
55     def postprocess_message(self, msg):
56         pass
57 
58     def __call__(self, *args, **kwargs):
59         with self:
60             return self.fn(*args, **kwargs)
61 
62 
63 def sample(name, fn, obs=None, sample_shape=()):
64     """
65     Returns a random sample from the stochastic function `fn`. This can have
66     additional side effects when wrapped inside effect handlers like
67     :class:`~numpyro.handlers.substitute`.
68 
69     :param str name: name of the sample site
70     :param fn: Python callable
71     :param numpy.ndarray obs: observed value
72     :param sample_shape: Shape of samples to be drawn.
73     :return: sample from the stochastic `fn`.
74     """
75     # if there are no active Messengers, we just draw a sample and return it as expected:
76     if not _PYRO_STACK:
77         return fn(sample_shape=sample_shape)
78 
79     # Otherwise, we initialize a message...
80     initial_msg = {
81         'type': 'sample',
82         'name': name,
83         'fn': fn,
84         'args': (),
85         'kwargs': {'sample_shape': sample_shape},
86         'value': obs,
87         'scale': 1.0,
88         'is_observed': obs is not None,
89         'intermediates': [],
90         'cond_indep_stack': [],
91     }
92 
93     # ...and use apply_stack to send it to the Messengers
94     msg = apply_stack(initial_msg)
95     return msg['value']
96 
97 
98 def identity(x, *args, **kwargs):
99     return x
100 
101 
102 def param(name, init_value=None, **kwargs):
103     """
104     Annotate the given site as an optimizable parameter for use with
105     :mod:`jax.experimental.optimizers`. For an example of how `param` statements
106     can be used in inference algorithms, refer to :func:`~numpyro.svi.svi`.
107 
108     :param str name: name of site.
109     :param numpy.ndarray init_value: initial value specified by the user. Note that
110         the onus of using this to initialize the optimizer is on the user /
111         inference algorithm, since there is no global parameter store in
112         NumPyro.
113     :return: value for the parameter. Unless wrapped inside a
114         handler like :class:`~numpyro.handlers.substitute`, this will simply
115         return the initial value.
116     """
117     # if there are no active Messengers, we just draw a sample and return it as expected:
118     if not _PYRO_STACK:
119         return init_value
120 
121     # Otherwise, we initialize a message...
122     initial_msg = {
123         'type': 'param',
124         'name': name,
125         'fn': identity,
126         'args': (init_value,),
127         'kwargs': kwargs,
128         'value': None,
129         'scale': 1.0,
130         'cond_indep_stack': [],
131     }
132 
133     # ...and use apply_stack to send it to the Messengers
134     msg = apply_stack(initial_msg)
135     return msg['value']
136 
137 
138 def module(name, nn, input_shape=None):
139     """
140     Declare a :mod:`~jax.experimental.stax` style neural network inside a
141     model so that its parameters are registered for optimization via
142     :func:`~numpyro.primitives.param` statements.
143 
144     :param str name: name of the module to be registered.
145     :param tuple nn: a tuple of `(init_fn, apply_fn)` obtained by a :mod:`~jax.experimental.stax`
146         constructor function.
147     :param tuple input_shape: shape of the input taken by the
148         neural network.
149     :return: a `apply_fn` with bound parameters that takes an array
150         as an input and returns the neural network transformed output
151         array.
152     """
153     module_key = name + '$params'
154     nn_init, nn_apply = nn
155     nn_params = param(module_key)
156     if nn_params is None:
157         if input_shape is None:
158             raise ValueError('Valid value for `input_size` needed to initialize.')
159         rng = numpyro.sample(name + '$rng', PRNGIdentity())
160         _, nn_params = nn_init(rng, input_shape)
161         param(module_key, nn_params)
162     return jax.partial(nn_apply, nn_params)
163 
164 
165 class plate(Messenger):
166     """
167     Construct for annotating conditionally independent variables. Within a
168     `plate` context manager, `sample` sites will be automatically broadcasted to
169     the size of the plate. Additionally, a scale factor might be applied by
170     certain inference algorithms if `subsample_size` is specified.
171 
172     :param str name: Name of the plate.
173     :param int size: Size of the plate.
174     :param int subsample_size: Optional argument denoting the size of the mini-batch.
175         This can be used to apply a scaling factor by inference algorithms. e.g.
176         when computing ELBO using a mini-batch.
177     :param int dim: Optional argument to specify which dimension in the tensor
178         is used as the plate dim. If `None` (default), the leftmost available dim
179         is allocated.
180     """
181     def __init__(self, name, size, subsample_size=None, dim=None):
182         self.name = name
183         self.size = size
184         self.subsample_size = size if subsample_size is None else subsample_size
185         if dim is not None and dim >= 0:
186             raise ValueError('dim arg must be negative.')
187         self.dim = dim
188         self._validate_and_set_dim()
189         super(plate, self).__init__()
190 
191     def _validate_and_set_dim(self):
192         msg = {
193             'type': 'plate',
194             'is_observed': False,
195             'fn': identity,
196             'name': self.name,
197             'args': (None,),
198             'kwargs': {},
199             'value': None,
200             'scale': 1.0,
201             'cond_indep_stack': [],
202         }
203         apply_stack(msg)
204         cond_indep_stack = msg['cond_indep_stack']
205         occupied_dims = {f.dim for f in cond_indep_stack}
206         dim = -1
207         while True:
208             if dim not in occupied_dims:
209                 break
210             dim -= 1
211         if self.dim is None:
212             self.dim = dim
213         else:
214             assert self.dim not in occupied_dims
215 
216     @staticmethod
217     def _get_batch_shape(cond_indep_stack):
218         n_dims = max(-f.dim for f in cond_indep_stack)
219         batch_shape = [1] * n_dims
220         for f in cond_indep_stack:
221             batch_shape[f.dim] = f.size
222         return tuple(batch_shape)
223 
224     def process_message(self, msg):
225         cond_indep_stack = msg['cond_indep_stack']
226         frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)
227         cond_indep_stack.append(frame)
228         batch_shape = self._get_batch_shape(cond_indep_stack)
229         if 'sample_shape' in msg['kwargs']:
230             batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)
231         msg['kwargs']['sample_shape'] = batch_shape
232         msg['scale'] = msg['scale'] * self.size / self.subsample_size
233
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -60,21 +60,28 @@
             return self.fn(*args, **kwargs)
 
 
-def sample(name, fn, obs=None, sample_shape=()):
+def sample(name, fn, obs=None, random_state=None, sample_shape=()):
     """
     Returns a random sample from the stochastic function `fn`. This can have
     additional side effects when wrapped inside effect handlers like
     :class:`~numpyro.handlers.substitute`.
 
+    .. note::
+        By design, `sample` primitive is meant to be used inside a NumPyro model.
+        Then :class:`~numpyro.handlers.seed` handler is used to inject a random
+        state to `fn`. In those situations, `random_state` keyword will take no
+        effect.
+
     :param str name: name of the sample site
     :param fn: Python callable
     :param numpy.ndarray obs: observed value
+    :param jax.random.PRNGKey random_state: an optional random key for `fn`.
     :param sample_shape: Shape of samples to be drawn.
     :return: sample from the stochastic `fn`.
     """
     # if there are no active Messengers, we just draw a sample and return it as expected:
     if not _PYRO_STACK:
-        return fn(sample_shape=sample_shape)
+        return fn(random_state=random_state, sample_shape=sample_shape)
 
     # Otherwise, we initialize a message...
     initial_msg = {
@@ -82,7 +89,7 @@
         'name': name,
         'fn': fn,
         'args': (),
-        'kwargs': {'sample_shape': sample_shape},
+        'kwargs': {'random_state': random_state, 'sample_shape': sample_shape},
         'value': obs,
         'scale': 1.0,
         'is_observed': obs is not None,

{"golden_diff": "diff --git a/numpyro/primitives.py b/numpyro/primitives.py\n--- a/numpyro/primitives.py\n+++ b/numpyro/primitives.py\n@@ -60,21 +60,28 @@\n             return self.fn(*args, **kwargs)\n \n \n-def sample(name, fn, obs=None, sample_shape=()):\n+def sample(name, fn, obs=None, random_state=None, sample_shape=()):\n     \"\"\"\n     Returns a random sample from the stochastic function `fn`. This can have\n     additional side effects when wrapped inside effect handlers like\n     :class:`~numpyro.handlers.substitute`.\n \n+    .. note::\n+        By design, `sample` primitive is meant to be used inside a NumPyro model.\n+        Then :class:`~numpyro.handlers.seed` handler is used to inject a random\n+        state to `fn`. In those situations, `random_state` keyword will take no\n+        effect.\n+\n     :param str name: name of the sample site\n     :param fn: Python callable\n     :param numpy.ndarray obs: observed value\n+    :param jax.random.PRNGKey random_state: an optional random key for `fn`.\n     :param sample_shape: Shape of samples to be drawn.\n     :return: sample from the stochastic `fn`.\n     \"\"\"\n     # if there are no active Messengers, we just draw a sample and return it as expected:\n     if not _PYRO_STACK:\n-        return fn(sample_shape=sample_shape)\n+        return fn(random_state=random_state, sample_shape=sample_shape)\n \n     # Otherwise, we initialize a message...\n     initial_msg = {\n@@ -82,7 +89,7 @@\n         'name': name,\n         'fn': fn,\n         'args': (),\n-        'kwargs': {'sample_shape': sample_shape},\n+        'kwargs': {'random_state': random_state, 'sample_shape': sample_shape},\n         'value': obs,\n         'scale': 1.0,\n         'is_observed': obs is not None,\n", "issue": "Document behavior of primitive statements in README\nHello! I'll preface this issue by stating that I'm new to numpyro, so there's a significant chance that I'm not using `numpyro.sample` correctly.\r\n\r\nWhen building models, I want to test draws from my random variables to make sure I'm passing parameters correctly, getting the kwargs correct, etc.; I'll lift a line out of a model function, [like this one](https://github.com/pyro-ppl/numpyro/blob/master/examples/baseball.py#L111), and run it in an IPython window to see if I get an error. It seems like I'm unable to do this by itself in an IPython console.\r\n\r\n## Minimal example\r\n\r\n```\r\nIn [1]: import numpyro\r\n\r\nIn [2]: import numpyro.distributions as d\r\n\r\nIn [3]: numpyro.__version__\r\nOut[3]: '0.2.0'\r\n\r\nIn [4]: numpyro.sample(\"x\", d.Normal(0, 1))\r\n---------------------------------------------------------------------------\r\nKeyError                                  Traceback (most recent call last)\r\n<ipython-input-4-21ce96e72ad0> in <module>\r\n----> 1 numpyro.sample(\"x\", d.Normal(0, 1))\r\n\r\n~/miniconda3/envs/fusion/lib/python3.7/site-packages/numpyro/primitives.py in sample(name, fn, obs, sample_shape)\r\n     45     # if there are no active Messengers, we just draw a sample and return it as expected:\r\n     46     if not _PYRO_STACK:\r\n---> 47         return fn(sample_shape=sample_shape)\r\n     48\r\n     49     # Otherwise, we initialize a message...\r\n\r\n~/miniconda3/envs/fusion/lib/python3.7/site-packages/numpyro/distributions/distribution.py in __call__(self, *args, **kwargs)\r\n    161\r\n    162     def __call__(self, *args, **kwargs):\r\n--> 163         key = kwargs.pop('random_state')\r\n    164         sample_intermediates = kwargs.pop('sample_intermediates', False)\r\n    165         if sample_intermediates:\r\n\r\nKeyError: 'random_state'\r\n```\r\n\r\n## Thoughts\r\nI can see that in [numpyro.sample](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/primitives.py#L47), we return `fn(sample_shape=sample_shape)`, which in this case attempts to call `dist.Normal(0, 1)(sample_shape=sample_shape)`. However, looking at [distributions.py](https://github.com/pyro-ppl/numpyro/blob/master/numpyro/distributions/distribution.py#L163), it seems that `Distribution.sample` expects a `random_state` kwarg that isn't getting passed.\r\n\r\nWhen I do the following, everything is fine:\r\n\r\n```\r\nIn [5]: import jax.random as random\r\n\r\nIn [6]: d.Normal(0, 1).sample(random.PRNGKey(12))\r\n/Users/chtu8001/miniconda3/envs/fusion/lib/python3.7/site-packages/jax/lib/xla_bridge.py:114: UserWarning: No GPU/TPU found, falling back to CPU.\r\n  warnings.warn('No GPU/TPU found, falling back to CPU.')\r\nOut[6]: DeviceArray(-0.5135873, dtype=float32)\r\n\r\nIn [7]: d.Normal(0, 1).sample(random.PRNGKey(12), sample_shape=(4, ))\r\nOut[7]:\r\nDeviceArray([-1.31179953, -0.70821768,  0.18907626, -1.09385514],\r\n            dtype=float32)\r\n```\r\n\r\nso I expect that we can't use `numpyro.sample` outside a model (although it's not totally clear to me how that is defined, something to do with Messengers and PYRO_STACK). I'm wondering if this is by design and I should just use the second, working example, or if I'm misunderstanding how `numpyro.sample` should be used.\r\n\r\nThanks!\n", "before_files": [{"content": "from collections import namedtuple\nimport functools\n\nimport jax\nfrom jax import lax\n\nimport numpyro\nfrom numpyro.distributions.discrete import PRNGIdentity\n\n_PYRO_STACK = []\n\n\nCondIndepStackFrame = namedtuple('CondIndepStackFrame', ['name', 'dim', 'size'])\n\n\ndef apply_stack(msg):\n    pointer = 0\n    for pointer, handler in enumerate(reversed(_PYRO_STACK)):\n        handler.process_message(msg)\n        # When a Messenger sets the \"stop\" field of a message,\n        # it prevents any Messengers above it on the stack from being applied.\n        if msg.get(\"stop\"):\n            break\n    if msg['value'] is None:\n        if msg['type'] == 'sample':\n            msg['value'], msg['intermediates'] = msg['fn'](*msg['args'],\n                                                           sample_intermediates=True,\n                                                           **msg['kwargs'])\n        else:\n            msg['value'] = msg['fn'](*msg['args'], **msg['kwargs'])\n\n    # A Messenger that sets msg[\"stop\"] == True also prevents application\n    # of postprocess_message by Messengers above it on the stack\n    # via the pointer variable from the process_message loop\n    for handler in _PYRO_STACK[-pointer-1:]:\n        handler.postprocess_message(msg)\n    return msg\n\n\nclass Messenger(object):\n    def __init__(self, fn=None):\n        self.fn = fn\n        functools.update_wrapper(self, fn, updated=[])\n\n    def __enter__(self):\n        _PYRO_STACK.append(self)\n\n    def __exit__(self, *args, **kwargs):\n        assert _PYRO_STACK[-1] is self\n        _PYRO_STACK.pop()\n\n    def process_message(self, msg):\n        pass\n\n    def postprocess_message(self, msg):\n        pass\n\n    def __call__(self, *args, **kwargs):\n        with self:\n            return self.fn(*args, **kwargs)\n\n\ndef sample(name, fn, obs=None, sample_shape=()):\n    \"\"\"\n    Returns a random sample from the stochastic function `fn`. This can have\n    additional side effects when wrapped inside effect handlers like\n    :class:`~numpyro.handlers.substitute`.\n\n    :param str name: name of the sample site\n    :param fn: Python callable\n    :param numpy.ndarray obs: observed value\n    :param sample_shape: Shape of samples to be drawn.\n    :return: sample from the stochastic `fn`.\n    \"\"\"\n    # if there are no active Messengers, we just draw a sample and return it as expected:\n    if not _PYRO_STACK:\n        return fn(sample_shape=sample_shape)\n\n    # Otherwise, we initialize a message...\n    initial_msg = {\n        'type': 'sample',\n        'name': name,\n        'fn': fn,\n        'args': (),\n        'kwargs': {'sample_shape': sample_shape},\n        'value': obs,\n        'scale': 1.0,\n        'is_observed': obs is not None,\n        'intermediates': [],\n        'cond_indep_stack': [],\n    }\n\n    # ...and use apply_stack to send it to the Messengers\n    msg = apply_stack(initial_msg)\n    return msg['value']\n\n\ndef identity(x, *args, **kwargs):\n    return x\n\n\ndef param(name, init_value=None, **kwargs):\n    \"\"\"\n    Annotate the given site as an optimizable parameter for use with\n    :mod:`jax.experimental.optimizers`. For an example of how `param` statements\n    can be used in inference algorithms, refer to :func:`~numpyro.svi.svi`.\n\n    :param str name: name of site.\n    :param numpy.ndarray init_value: initial value specified by the user. Note that\n        the onus of using this to initialize the optimizer is on the user /\n        inference algorithm, since there is no global parameter store in\n        NumPyro.\n    :return: value for the parameter. Unless wrapped inside a\n        handler like :class:`~numpyro.handlers.substitute`, this will simply\n        return the initial value.\n    \"\"\"\n    # if there are no active Messengers, we just draw a sample and return it as expected:\n    if not _PYRO_STACK:\n        return init_value\n\n    # Otherwise, we initialize a message...\n    initial_msg = {\n        'type': 'param',\n        'name': name,\n        'fn': identity,\n        'args': (init_value,),\n        'kwargs': kwargs,\n        'value': None,\n        'scale': 1.0,\n        'cond_indep_stack': [],\n    }\n\n    # ...and use apply_stack to send it to the Messengers\n    msg = apply_stack(initial_msg)\n    return msg['value']\n\n\ndef module(name, nn, input_shape=None):\n    \"\"\"\n    Declare a :mod:`~jax.experimental.stax` style neural network inside a\n    model so that its parameters are registered for optimization via\n    :func:`~numpyro.primitives.param` statements.\n\n    :param str name: name of the module to be registered.\n    :param tuple nn: a tuple of `(init_fn, apply_fn)` obtained by a :mod:`~jax.experimental.stax`\n        constructor function.\n    :param tuple input_shape: shape of the input taken by the\n        neural network.\n    :return: a `apply_fn` with bound parameters that takes an array\n        as an input and returns the neural network transformed output\n        array.\n    \"\"\"\n    module_key = name + '$params'\n    nn_init, nn_apply = nn\n    nn_params = param(module_key)\n    if nn_params is None:\n        if input_shape is None:\n            raise ValueError('Valid value for `input_size` needed to initialize.')\n        rng = numpyro.sample(name + '$rng', PRNGIdentity())\n        _, nn_params = nn_init(rng, input_shape)\n        param(module_key, nn_params)\n    return jax.partial(nn_apply, nn_params)\n\n\nclass plate(Messenger):\n    \"\"\"\n    Construct for annotating conditionally independent variables. Within a\n    `plate` context manager, `sample` sites will be automatically broadcasted to\n    the size of the plate. Additionally, a scale factor might be applied by\n    certain inference algorithms if `subsample_size` is specified.\n\n    :param str name: Name of the plate.\n    :param int size: Size of the plate.\n    :param int subsample_size: Optional argument denoting the size of the mini-batch.\n        This can be used to apply a scaling factor by inference algorithms. e.g.\n        when computing ELBO using a mini-batch.\n    :param int dim: Optional argument to specify which dimension in the tensor\n        is used as the plate dim. If `None` (default), the leftmost available dim\n        is allocated.\n    \"\"\"\n    def __init__(self, name, size, subsample_size=None, dim=None):\n        self.name = name\n        self.size = size\n        self.subsample_size = size if subsample_size is None else subsample_size\n        if dim is not None and dim >= 0:\n            raise ValueError('dim arg must be negative.')\n        self.dim = dim\n        self._validate_and_set_dim()\n        super(plate, self).__init__()\n\n    def _validate_and_set_dim(self):\n        msg = {\n            'type': 'plate',\n            'is_observed': False,\n            'fn': identity,\n            'name': self.name,\n            'args': (None,),\n            'kwargs': {},\n            'value': None,\n            'scale': 1.0,\n            'cond_indep_stack': [],\n        }\n        apply_stack(msg)\n        cond_indep_stack = msg['cond_indep_stack']\n        occupied_dims = {f.dim for f in cond_indep_stack}\n        dim = -1\n        while True:\n            if dim not in occupied_dims:\n                break\n            dim -= 1\n        if self.dim is None:\n            self.dim = dim\n        else:\n            assert self.dim not in occupied_dims\n\n    @staticmethod\n    def _get_batch_shape(cond_indep_stack):\n        n_dims = max(-f.dim for f in cond_indep_stack)\n        batch_shape = [1] * n_dims\n        for f in cond_indep_stack:\n            batch_shape[f.dim] = f.size\n        return tuple(batch_shape)\n\n    def process_message(self, msg):\n        cond_indep_stack = msg['cond_indep_stack']\n        frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)\n        cond_indep_stack.append(frame)\n        batch_shape = self._get_batch_shape(cond_indep_stack)\n        if 'sample_shape' in msg['kwargs']:\n            batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)\n        msg['kwargs']['sample_shape'] = batch_shape\n        msg['scale'] = msg['scale'] * self.size / self.subsample_size\n", "path": "numpyro/primitives.py"}], "after_files": [{"content": "from collections import namedtuple\nimport functools\n\nimport jax\nfrom jax import lax\n\nimport numpyro\nfrom numpyro.distributions.discrete import PRNGIdentity\n\n_PYRO_STACK = []\n\n\nCondIndepStackFrame = namedtuple('CondIndepStackFrame', ['name', 'dim', 'size'])\n\n\ndef apply_stack(msg):\n    pointer = 0\n    for pointer, handler in enumerate(reversed(_PYRO_STACK)):\n        handler.process_message(msg)\n        # When a Messenger sets the \"stop\" field of a message,\n        # it prevents any Messengers above it on the stack from being applied.\n        if msg.get(\"stop\"):\n            break\n    if msg['value'] is None:\n        if msg['type'] == 'sample':\n            msg['value'], msg['intermediates'] = msg['fn'](*msg['args'],\n                                                           sample_intermediates=True,\n                                                           **msg['kwargs'])\n        else:\n            msg['value'] = msg['fn'](*msg['args'], **msg['kwargs'])\n\n    # A Messenger that sets msg[\"stop\"] == True also prevents application\n    # of postprocess_message by Messengers above it on the stack\n    # via the pointer variable from the process_message loop\n    for handler in _PYRO_STACK[-pointer-1:]:\n        handler.postprocess_message(msg)\n    return msg\n\n\nclass Messenger(object):\n    def __init__(self, fn=None):\n        self.fn = fn\n        functools.update_wrapper(self, fn, updated=[])\n\n    def __enter__(self):\n        _PYRO_STACK.append(self)\n\n    def __exit__(self, *args, **kwargs):\n        assert _PYRO_STACK[-1] is self\n        _PYRO_STACK.pop()\n\n    def process_message(self, msg):\n        pass\n\n    def postprocess_message(self, msg):\n        pass\n\n    def __call__(self, *args, **kwargs):\n        with self:\n            return self.fn(*args, **kwargs)\n\n\ndef sample(name, fn, obs=None, random_state=None, sample_shape=()):\n    \"\"\"\n    Returns a random sample from the stochastic function `fn`. This can have\n    additional side effects when wrapped inside effect handlers like\n    :class:`~numpyro.handlers.substitute`.\n\n    .. note::\n        By design, `sample` primitive is meant to be used inside a NumPyro model.\n        Then :class:`~numpyro.handlers.seed` handler is used to inject a random\n        state to `fn`. In those situations, `random_state` keyword will take no\n        effect.\n\n    :param str name: name of the sample site\n    :param fn: Python callable\n    :param numpy.ndarray obs: observed value\n    :param jax.random.PRNGKey random_state: an optional random key for `fn`.\n    :param sample_shape: Shape of samples to be drawn.\n    :return: sample from the stochastic `fn`.\n    \"\"\"\n    # if there are no active Messengers, we just draw a sample and return it as expected:\n    if not _PYRO_STACK:\n        return fn(random_state=random_state, sample_shape=sample_shape)\n\n    # Otherwise, we initialize a message...\n    initial_msg = {\n        'type': 'sample',\n        'name': name,\n        'fn': fn,\n        'args': (),\n        'kwargs': {'random_state': random_state, 'sample_shape': sample_shape},\n        'value': obs,\n        'scale': 1.0,\n        'is_observed': obs is not None,\n        'intermediates': [],\n        'cond_indep_stack': [],\n    }\n\n    # ...and use apply_stack to send it to the Messengers\n    msg = apply_stack(initial_msg)\n    return msg['value']\n\n\ndef identity(x, *args, **kwargs):\n    return x\n\n\ndef param(name, init_value=None, **kwargs):\n    \"\"\"\n    Annotate the given site as an optimizable parameter for use with\n    :mod:`jax.experimental.optimizers`. For an example of how `param` statements\n    can be used in inference algorithms, refer to :func:`~numpyro.svi.svi`.\n\n    :param str name: name of site.\n    :param numpy.ndarray init_value: initial value specified by the user. Note that\n        the onus of using this to initialize the optimizer is on the user /\n        inference algorithm, since there is no global parameter store in\n        NumPyro.\n    :return: value for the parameter. Unless wrapped inside a\n        handler like :class:`~numpyro.handlers.substitute`, this will simply\n        return the initial value.\n    \"\"\"\n    # if there are no active Messengers, we just draw a sample and return it as expected:\n    if not _PYRO_STACK:\n        return init_value\n\n    # Otherwise, we initialize a message...\n    initial_msg = {\n        'type': 'param',\n        'name': name,\n        'fn': identity,\n        'args': (init_value,),\n        'kwargs': kwargs,\n        'value': None,\n        'scale': 1.0,\n        'cond_indep_stack': [],\n    }\n\n    # ...and use apply_stack to send it to the Messengers\n    msg = apply_stack(initial_msg)\n    return msg['value']\n\n\ndef module(name, nn, input_shape=None):\n    \"\"\"\n    Declare a :mod:`~jax.experimental.stax` style neural network inside a\n    model so that its parameters are registered for optimization via\n    :func:`~numpyro.primitives.param` statements.\n\n    :param str name: name of the module to be registered.\n    :param tuple nn: a tuple of `(init_fn, apply_fn)` obtained by a :mod:`~jax.experimental.stax`\n        constructor function.\n    :param tuple input_shape: shape of the input taken by the\n        neural network.\n    :return: a `apply_fn` with bound parameters that takes an array\n        as an input and returns the neural network transformed output\n        array.\n    \"\"\"\n    module_key = name + '$params'\n    nn_init, nn_apply = nn\n    nn_params = param(module_key)\n    if nn_params is None:\n        if input_shape is None:\n            raise ValueError('Valid value for `input_size` needed to initialize.')\n        rng = numpyro.sample(name + '$rng', PRNGIdentity())\n        _, nn_params = nn_init(rng, input_shape)\n        param(module_key, nn_params)\n    return jax.partial(nn_apply, nn_params)\n\n\nclass plate(Messenger):\n    \"\"\"\n    Construct for annotating conditionally independent variables. Within a\n    `plate` context manager, `sample` sites will be automatically broadcasted to\n    the size of the plate. Additionally, a scale factor might be applied by\n    certain inference algorithms if `subsample_size` is specified.\n\n    :param str name: Name of the plate.\n    :param int size: Size of the plate.\n    :param int subsample_size: Optional argument denoting the size of the mini-batch.\n        This can be used to apply a scaling factor by inference algorithms. e.g.\n        when computing ELBO using a mini-batch.\n    :param int dim: Optional argument to specify which dimension in the tensor\n        is used as the plate dim. If `None` (default), the leftmost available dim\n        is allocated.\n    \"\"\"\n    def __init__(self, name, size, subsample_size=None, dim=None):\n        self.name = name\n        self.size = size\n        self.subsample_size = size if subsample_size is None else subsample_size\n        if dim is not None and dim >= 0:\n            raise ValueError('dim arg must be negative.')\n        self.dim = dim\n        self._validate_and_set_dim()\n        super(plate, self).__init__()\n\n    def _validate_and_set_dim(self):\n        msg = {\n            'type': 'plate',\n            'is_observed': False,\n            'fn': identity,\n            'name': self.name,\n            'args': (None,),\n            'kwargs': {},\n            'value': None,\n            'scale': 1.0,\n            'cond_indep_stack': [],\n        }\n        apply_stack(msg)\n        cond_indep_stack = msg['cond_indep_stack']\n        occupied_dims = {f.dim for f in cond_indep_stack}\n        dim = -1\n        while True:\n            if dim not in occupied_dims:\n                break\n            dim -= 1\n        if self.dim is None:\n            self.dim = dim\n        else:\n            assert self.dim not in occupied_dims\n\n    @staticmethod\n    def _get_batch_shape(cond_indep_stack):\n        n_dims = max(-f.dim for f in cond_indep_stack)\n        batch_shape = [1] * n_dims\n        for f in cond_indep_stack:\n            batch_shape[f.dim] = f.size\n        return tuple(batch_shape)\n\n    def process_message(self, msg):\n        cond_indep_stack = msg['cond_indep_stack']\n        frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)\n        cond_indep_stack.append(frame)\n        batch_shape = self._get_batch_shape(cond_indep_stack)\n        if 'sample_shape' in msg['kwargs']:\n            batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)\n        msg['kwargs']['sample_shape'] = batch_shape\n        msg['scale'] = msg['scale'] * self.size / self.subsample_size\n", "path": "numpyro/primitives.py"}]}