We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
`plate` has unexpected behaviour when dist has batch_shape
I faced this issue while revising baseball example using `plate`. The repro code
```
import jax.numpy as np
import numpyro
import numpyro.distributions as dist
from numpyro.handlers import trace, seed

def model():
    with numpyro.plate('plate', 10):
        numpyro.sample('x', dist.Normal(np.zeros(10), 1))

trace(seed(model, 0)).get_trace()['x']['value'].shape
```
should return `(10,)` but currently, it returns `(10, 10)`.

cc @neerajprad 

--- 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, random_state=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     .. note::
70         By design, `sample` primitive is meant to be used inside a NumPyro model.
71         Then :class:`~numpyro.handlers.seed` handler is used to inject a random
72         state to `fn`. In those situations, `random_state` keyword will take no
73         effect.
74 
75     :param str name: name of the sample site
76     :param fn: Python callable
77     :param numpy.ndarray obs: observed value
78     :param jax.random.PRNGKey random_state: an optional random key for `fn`.
79     :param sample_shape: Shape of samples to be drawn.
80     :return: sample from the stochastic `fn`.
81     """
82     # if there are no active Messengers, we just draw a sample and return it as expected:
83     if not _PYRO_STACK:
84         return fn(random_state=random_state, sample_shape=sample_shape)
85 
86     # Otherwise, we initialize a message...
87     initial_msg = {
88         'type': 'sample',
89         'name': name,
90         'fn': fn,
91         'args': (),
92         'kwargs': {'random_state': random_state, 'sample_shape': sample_shape},
93         'value': obs,
94         'scale': 1.0,
95         'is_observed': obs is not None,
96         'intermediates': [],
97         'cond_indep_stack': [],
98     }
99 
100     # ...and use apply_stack to send it to the Messengers
101     msg = apply_stack(initial_msg)
102     return msg['value']
103 
104 
105 def identity(x, *args, **kwargs):
106     return x
107 
108 
109 def param(name, init_value=None, **kwargs):
110     """
111     Annotate the given site as an optimizable parameter for use with
112     :mod:`jax.experimental.optimizers`. For an example of how `param` statements
113     can be used in inference algorithms, refer to :func:`~numpyro.svi.svi`.
114 
115     :param str name: name of site.
116     :param numpy.ndarray init_value: initial value specified by the user. Note that
117         the onus of using this to initialize the optimizer is on the user /
118         inference algorithm, since there is no global parameter store in
119         NumPyro.
120     :return: value for the parameter. Unless wrapped inside a
121         handler like :class:`~numpyro.handlers.substitute`, this will simply
122         return the initial value.
123     """
124     # if there are no active Messengers, we just draw a sample and return it as expected:
125     if not _PYRO_STACK:
126         return init_value
127 
128     # Otherwise, we initialize a message...
129     initial_msg = {
130         'type': 'param',
131         'name': name,
132         'fn': identity,
133         'args': (init_value,),
134         'kwargs': kwargs,
135         'value': None,
136         'scale': 1.0,
137         'cond_indep_stack': [],
138     }
139 
140     # ...and use apply_stack to send it to the Messengers
141     msg = apply_stack(initial_msg)
142     return msg['value']
143 
144 
145 def module(name, nn, input_shape=None):
146     """
147     Declare a :mod:`~jax.experimental.stax` style neural network inside a
148     model so that its parameters are registered for optimization via
149     :func:`~numpyro.primitives.param` statements.
150 
151     :param str name: name of the module to be registered.
152     :param tuple nn: a tuple of `(init_fn, apply_fn)` obtained by a :mod:`~jax.experimental.stax`
153         constructor function.
154     :param tuple input_shape: shape of the input taken by the
155         neural network.
156     :return: a `apply_fn` with bound parameters that takes an array
157         as an input and returns the neural network transformed output
158         array.
159     """
160     module_key = name + '$params'
161     nn_init, nn_apply = nn
162     nn_params = param(module_key)
163     if nn_params is None:
164         if input_shape is None:
165             raise ValueError('Valid value for `input_size` needed to initialize.')
166         rng = numpyro.sample(name + '$rng', PRNGIdentity())
167         _, nn_params = nn_init(rng, input_shape)
168         param(module_key, nn_params)
169     return jax.partial(nn_apply, nn_params)
170 
171 
172 class plate(Messenger):
173     """
174     Construct for annotating conditionally independent variables. Within a
175     `plate` context manager, `sample` sites will be automatically broadcasted to
176     the size of the plate. Additionally, a scale factor might be applied by
177     certain inference algorithms if `subsample_size` is specified.
178 
179     :param str name: Name of the plate.
180     :param int size: Size of the plate.
181     :param int subsample_size: Optional argument denoting the size of the mini-batch.
182         This can be used to apply a scaling factor by inference algorithms. e.g.
183         when computing ELBO using a mini-batch.
184     :param int dim: Optional argument to specify which dimension in the tensor
185         is used as the plate dim. If `None` (default), the leftmost available dim
186         is allocated.
187     """
188     def __init__(self, name, size, subsample_size=None, dim=None):
189         self.name = name
190         self.size = size
191         self.subsample_size = size if subsample_size is None else subsample_size
192         if dim is not None and dim >= 0:
193             raise ValueError('dim arg must be negative.')
194         self.dim = dim
195         self._validate_and_set_dim()
196         super(plate, self).__init__()
197 
198     def _validate_and_set_dim(self):
199         msg = {
200             'type': 'plate',
201             'fn': identity,
202             'name': self.name,
203             'args': (None,),
204             'kwargs': {},
205             'value': None,
206             'scale': 1.0,
207             'cond_indep_stack': [],
208         }
209         apply_stack(msg)
210         cond_indep_stack = msg['cond_indep_stack']
211         occupied_dims = {f.dim for f in cond_indep_stack}
212         dim = -1
213         while True:
214             if dim not in occupied_dims:
215                 break
216             dim -= 1
217         if self.dim is None:
218             self.dim = dim
219         else:
220             assert self.dim not in occupied_dims
221 
222     @staticmethod
223     def _get_batch_shape(cond_indep_stack):
224         n_dims = max(-f.dim for f in cond_indep_stack)
225         batch_shape = [1] * n_dims
226         for f in cond_indep_stack:
227             batch_shape[f.dim] = f.size
228         return tuple(batch_shape)
229 
230     def process_message(self, msg):
231         cond_indep_stack = msg['cond_indep_stack']
232         frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)
233         cond_indep_stack.append(frame)
234         batch_shape = self._get_batch_shape(cond_indep_stack)
235         if 'sample_shape' in msg['kwargs']:
236             batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)
237         msg['kwargs']['sample_shape'] = batch_shape
238         msg['scale'] = msg['scale'] * self.size / self.subsample_size
239
```
--- 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
@@ -231,7 +231,20 @@
         cond_indep_stack = msg['cond_indep_stack']
         frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)
         cond_indep_stack.append(frame)
-        batch_shape = self._get_batch_shape(cond_indep_stack)
+        expected_shape = self._get_batch_shape(cond_indep_stack)
+        dist_batch_shape = msg['fn'].batch_shape if msg['type'] == 'sample' else ()
+        overlap_idx = len(expected_shape) - len(dist_batch_shape)
+        if overlap_idx < 0:
+            raise ValueError('Expected dimensions within plate = {}, which is less than the '
+                             'distribution\'s batch shape = {}.'.format(len(expected_shape), len(dist_batch_shape)))
+        trailing_shape = expected_shape[overlap_idx:]
+        # e.g. distribution with batch shape (1, 5) cannot be broadcast to (5, 5)
+        broadcast_shape = lax.broadcast_shapes(trailing_shape, dist_batch_shape)
+        if broadcast_shape != dist_batch_shape:
+            raise ValueError('Distribution batch shape = {} cannot be broadcast up to {}. '
+                             'Consider using unbatched distributions.'
+                             .format(dist_batch_shape, broadcast_shape))
+        batch_shape = expected_shape[:overlap_idx]
         if 'sample_shape' in msg['kwargs']:
             batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)
         msg['kwargs']['sample_shape'] = batch_shape

{"golden_diff": "diff --git a/numpyro/primitives.py b/numpyro/primitives.py\n--- a/numpyro/primitives.py\n+++ b/numpyro/primitives.py\n@@ -231,7 +231,20 @@\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+        expected_shape = self._get_batch_shape(cond_indep_stack)\n+        dist_batch_shape = msg['fn'].batch_shape if msg['type'] == 'sample' else ()\n+        overlap_idx = len(expected_shape) - len(dist_batch_shape)\n+        if overlap_idx < 0:\n+            raise ValueError('Expected dimensions within plate = {}, which is less than the '\n+                             'distribution\\'s batch shape = {}.'.format(len(expected_shape), len(dist_batch_shape)))\n+        trailing_shape = expected_shape[overlap_idx:]\n+        # e.g. distribution with batch shape (1, 5) cannot be broadcast to (5, 5)\n+        broadcast_shape = lax.broadcast_shapes(trailing_shape, dist_batch_shape)\n+        if broadcast_shape != dist_batch_shape:\n+            raise ValueError('Distribution batch shape = {} cannot be broadcast up to {}. '\n+                             'Consider using unbatched distributions.'\n+                             .format(dist_batch_shape, broadcast_shape))\n+        batch_shape = expected_shape[:overlap_idx]\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", "issue": "`plate` has unexpected behaviour when dist has batch_shape\nI faced this issue while revising baseball example using `plate`. The repro code\r\n```\r\nimport jax.numpy as np\r\nimport numpyro\r\nimport numpyro.distributions as dist\r\nfrom numpyro.handlers import trace, seed\r\n\r\ndef model():\r\n    with numpyro.plate('plate', 10):\r\n        numpyro.sample('x', dist.Normal(np.zeros(10), 1))\r\n\r\ntrace(seed(model, 0)).get_trace()['x']['value'].shape\r\n```\r\nshould return `(10,)` but currently, it returns `(10, 10)`.\r\n\r\ncc @neerajprad \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, 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            '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            '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        expected_shape = self._get_batch_shape(cond_indep_stack)\n        dist_batch_shape = msg['fn'].batch_shape if msg['type'] == 'sample' else ()\n        overlap_idx = len(expected_shape) - len(dist_batch_shape)\n        if overlap_idx < 0:\n            raise ValueError('Expected dimensions within plate = {}, which is less than the '\n                             'distribution\\'s batch shape = {}.'.format(len(expected_shape), len(dist_batch_shape)))\n        trailing_shape = expected_shape[overlap_idx:]\n        # e.g. distribution with batch shape (1, 5) cannot be broadcast to (5, 5)\n        broadcast_shape = lax.broadcast_shapes(trailing_shape, dist_batch_shape)\n        if broadcast_shape != dist_batch_shape:\n            raise ValueError('Distribution batch shape = {} cannot be broadcast up to {}. '\n                             'Consider using unbatched distributions.'\n                             .format(dist_batch_shape, broadcast_shape))\n        batch_shape = expected_shape[:overlap_idx]\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"}]}