We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
ZeroDivisionError when running SVI with num_steps < 20
Reproducible code
```python
import jax
import numpyro

def model():
    pass

def guide():
    pass

svi = numpyro.infer.SVI(model, guide, numpyro.optim.Adam(1), numpyro.infer.Trace_ELBO())
svi.run(jax.random.PRNGKey(0), 10)
```

Proposed fix: setting `batch = max(num_steps // 20, 1)` in [SVI.run](https://github.com/pyro-ppl/numpyro/blob/aee945895dff9ffb6a12637a9bfbddb1d8c20363/numpyro/infer/svi.py#L179) to guarantee that `batch` is positive

--- 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/infer/svi.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 from functools import namedtuple, partial
5 
6 import tqdm
7 
8 from jax import jit, lax, random
9 import jax.numpy as jnp
10 
11 from numpyro.distributions import constraints
12 from numpyro.distributions.transforms import biject_to
13 from numpyro.handlers import replay, seed, trace
14 from numpyro.infer.util import transform_fn
15 
16 SVIState = namedtuple('SVIState', ['optim_state', 'rng_key'])
17 """
18 A :func:`~collections.namedtuple` consisting of the following fields:
19  - **optim_state** - current optimizer's state.
20  - **rng_key** - random number generator seed used for the iteration.
21 """
22 
23 
24 SVIRunResult = namedtuple('SVIRunResult', ['params', 'losses'])
25 """
26 A :func:`~collections.namedtuple` consisting of the following fields:
27  - **params** - the optimized parameters.
28  - **losses** - the losses collected at every step.
29 """
30 
31 
32 def _apply_loss_fn(loss_fn, rng_key, constrain_fn, model, guide,
33                    args, kwargs, static_kwargs, params):
34     return loss_fn(rng_key, constrain_fn(params), model, guide, *args, **kwargs, **static_kwargs)
35 
36 
37 class SVI(object):
38     """
39     Stochastic Variational Inference given an ELBO loss objective.
40 
41     **References**
42 
43     1. *SVI Part I: An Introduction to Stochastic Variational Inference in Pyro*,
44        (http://pyro.ai/examples/svi_part_i.html)
45 
46     **Example:**
47 
48     .. doctest::
49 
50         >>> from jax import random
51         >>> import jax.numpy as jnp
52         >>> import numpyro
53         >>> import numpyro.distributions as dist
54         >>> from numpyro.distributions import constraints
55         >>> from numpyro.infer import SVI, Trace_ELBO
56 
57         >>> def model(data):
58         ...     f = numpyro.sample("latent_fairness", dist.Beta(10, 10))
59         ...     with numpyro.plate("N", data.shape[0]):
60         ...         numpyro.sample("obs", dist.Bernoulli(f), obs=data)
61 
62         >>> def guide(data):
63         ...     alpha_q = numpyro.param("alpha_q", 15., constraint=constraints.positive)
64         ...     beta_q = numpyro.param("beta_q", lambda rng_key: random.exponential(rng_key),
65         ...                            constraint=constraints.positive)
66         ...     numpyro.sample("latent_fairness", dist.Beta(alpha_q, beta_q))
67 
68         >>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])
69         >>> optimizer = numpyro.optim.Adam(step_size=0.0005)
70         >>> svi = SVI(model, guide, optimizer, loss=Trace_ELBO())
71         >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)
72         >>> params = svi_result.params
73         >>> inferred_mean = params["alpha_q"] / (params["alpha_q"] + params["beta_q"])
74 
75     :param model: Python callable with Pyro primitives for the model.
76     :param guide: Python callable with Pyro primitives for the guide
77         (recognition network).
78     :param optim: an instance of :class:`~numpyro.optim._NumpyroOptim`.
79     :param loss: ELBO loss, i.e. negative Evidence Lower Bound, to minimize.
80     :param static_kwargs: static arguments for the model / guide, i.e. arguments
81         that remain constant during fitting.
82     :return: tuple of `(init_fn, update_fn, evaluate)`.
83     """
84     def __init__(self, model, guide, optim, loss, **static_kwargs):
85         self.model = model
86         self.guide = guide
87         self.loss = loss
88         self.optim = optim
89         self.static_kwargs = static_kwargs
90         self.constrain_fn = None
91 
92     def init(self, rng_key, *args, **kwargs):
93         """
94         Gets the initial SVI state.
95 
96         :param jax.random.PRNGKey rng_key: random number generator seed.
97         :param args: arguments to the model / guide (these can possibly vary during
98             the course of fitting).
99         :param kwargs: keyword arguments to the model / guide (these can possibly vary
100             during the course of fitting).
101         :return: the initial :data:`SVIState`
102         """
103         rng_key, model_seed, guide_seed = random.split(rng_key, 3)
104         model_init = seed(self.model, model_seed)
105         guide_init = seed(self.guide, guide_seed)
106         guide_trace = trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)
107         model_trace = trace(replay(model_init, guide_trace)).get_trace(*args, **kwargs, **self.static_kwargs)
108         params = {}
109         inv_transforms = {}
110         # NB: params in model_trace will be overwritten by params in guide_trace
111         for site in list(model_trace.values()) + list(guide_trace.values()):
112             if site['type'] == 'param':
113                 constraint = site['kwargs'].pop('constraint', constraints.real)
114                 transform = biject_to(constraint)
115                 inv_transforms[site['name']] = transform
116                 params[site['name']] = transform.inv(site['value'])
117 
118         self.constrain_fn = partial(transform_fn, inv_transforms)
119         return SVIState(self.optim.init(params), rng_key)
120 
121     def get_params(self, svi_state):
122         """
123         Gets values at `param` sites of the `model` and `guide`.
124 
125         :param svi_state: current state of SVI.
126         :return: the corresponding parameters
127         """
128         params = self.constrain_fn(self.optim.get_params(svi_state.optim_state))
129         return params
130 
131     def update(self, svi_state, *args, **kwargs):
132         """
133         Take a single step of SVI (possibly on a batch / minibatch of data),
134         using the optimizer.
135 
136         :param svi_state: current state of SVI.
137         :param args: arguments to the model / guide (these can possibly vary during
138             the course of fitting).
139         :param kwargs: keyword arguments to the model / guide (these can possibly vary
140             during the course of fitting).
141         :return: tuple of `(svi_state, loss)`.
142         """
143         rng_key, rng_key_step = random.split(svi_state.rng_key)
144         loss_fn = partial(_apply_loss_fn, self.loss.loss, rng_key_step, self.constrain_fn, self.model,
145                           self.guide, args, kwargs, self.static_kwargs)
146         loss_val, optim_state = self.optim.eval_and_update(loss_fn, svi_state.optim_state)
147         return SVIState(optim_state, rng_key), loss_val
148 
149     def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):
150         """
151         (EXPERIMENTAL INTERFACE) Run SVI with `num_steps` iterations, then return
152         the optimized parameters and the stacked losses at every step. If `num_steps`
153         is large, setting `progress_bar=False` can make the run faster.
154 
155         .. note:: For a complex training process (e.g. the one requires early stopping,
156             epoch training, varying args/kwargs,...), we recommend to use the more
157             flexible methods :meth:`init`, :meth:`update`, :meth:`evaluate` to
158             customize your training procedure.
159 
160         :param jax.random.PRNGKey rng_key: random number generator seed.
161         :param int num_steps: the number of optimization steps.
162         :param args: arguments to the model / guide
163         :param bool progress_bar: Whether to enable progress bar updates. Defaults to
164             ``True``.
165         :param kwargs: keyword arguments to the model / guide
166         :return: a namedtuple with fields `params` and `losses` where `params`
167             holds the optimized values at :class:`numpyro.param` sites,
168             and `losses` is the collected loss during the process.
169         :rtype: SVIRunResult
170         """
171         def body_fn(svi_state, carry):
172             svi_state, loss = self.update(svi_state, *args, **kwargs)
173             return svi_state, loss
174 
175         svi_state = self.init(rng_key, *args, **kwargs)
176         if progress_bar:
177             losses = []
178             with tqdm.trange(1, num_steps + 1) as t:
179                 batch = num_steps // 20
180                 for i in t:
181                     svi_state, loss = jit(body_fn)(svi_state, None)
182                     if i % batch == 0:
183                         avg_loss = sum(losses[i-batch:]) / batch
184                         t.set_postfix_str("init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}"
185                                           .format(losses[0], i - batch + 1, i, avg_loss),
186                                           refresh=False)
187                     losses.append(loss)
188             losses = jnp.stack(losses)
189         else:
190             svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)
191 
192         return SVIRunResult(self.get_params(svi_state), losses)
193 
194     def evaluate(self, svi_state, *args, **kwargs):
195         """
196         Take a single step of SVI (possibly on a batch / minibatch of data).
197 
198         :param svi_state: current state of SVI.
199         :param args: arguments to the model / guide (these can possibly vary during
200             the course of fitting).
201         :param kwargs: keyword arguments to the model / guide.
202         :return: evaluate ELBO loss given the current parameter values
203             (held within `svi_state.optim_state`).
204         """
205         # we split to have the same seed as `update_fn` given an svi_state
206         _, rng_key_eval = random.split(svi_state.rng_key)
207         params = self.get_params(svi_state)
208         return self.loss.loss(rng_key_eval, params, self.model, self.guide,
209                               *args, **kwargs, **self.static_kwargs)
210
```
--- 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/infer/svi.py b/numpyro/infer/svi.py
--- a/numpyro/infer/svi.py
+++ b/numpyro/infer/svi.py
@@ -176,15 +176,15 @@
         if progress_bar:
             losses = []
             with tqdm.trange(1, num_steps + 1) as t:
-                batch = num_steps // 20
+                batch = max(num_steps // 20, 1)
                 for i in t:
                     svi_state, loss = jit(body_fn)(svi_state, None)
+                    losses.append(loss)
                     if i % batch == 0:
                         avg_loss = sum(losses[i-batch:]) / batch
                         t.set_postfix_str("init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}"
                                           .format(losses[0], i - batch + 1, i, avg_loss),
                                           refresh=False)
-                    losses.append(loss)
             losses = jnp.stack(losses)
         else:
             svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)

{"golden_diff": "diff --git a/numpyro/infer/svi.py b/numpyro/infer/svi.py\n--- a/numpyro/infer/svi.py\n+++ b/numpyro/infer/svi.py\n@@ -176,15 +176,15 @@\n         if progress_bar:\n             losses = []\n             with tqdm.trange(1, num_steps + 1) as t:\n-                batch = num_steps // 20\n+                batch = max(num_steps // 20, 1)\n                 for i in t:\n                     svi_state, loss = jit(body_fn)(svi_state, None)\n+                    losses.append(loss)\n                     if i % batch == 0:\n                         avg_loss = sum(losses[i-batch:]) / batch\n                         t.set_postfix_str(\"init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}\"\n                                           .format(losses[0], i - batch + 1, i, avg_loss),\n                                           refresh=False)\n-                    losses.append(loss)\n             losses = jnp.stack(losses)\n         else:\n             svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)\n", "issue": "ZeroDivisionError when running SVI with num_steps < 20\nReproducible code\r\n```python\r\nimport jax\r\nimport numpyro\r\n\r\ndef model():\r\n    pass\r\n\r\ndef guide():\r\n    pass\r\n\r\nsvi = numpyro.infer.SVI(model, guide, numpyro.optim.Adam(1), numpyro.infer.Trace_ELBO())\r\nsvi.run(jax.random.PRNGKey(0), 10)\r\n```\r\n\r\nProposed fix: setting `batch = max(num_steps // 20, 1)` in [SVI.run](https://github.com/pyro-ppl/numpyro/blob/aee945895dff9ffb6a12637a9bfbddb1d8c20363/numpyro/infer/svi.py#L179) to guarantee that `batch` is positive\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom functools import namedtuple, partial\n\nimport tqdm\n\nfrom jax import jit, lax, random\nimport jax.numpy as jnp\n\nfrom numpyro.distributions import constraints\nfrom numpyro.distributions.transforms import biject_to\nfrom numpyro.handlers import replay, seed, trace\nfrom numpyro.infer.util import transform_fn\n\nSVIState = namedtuple('SVIState', ['optim_state', 'rng_key'])\n\"\"\"\nA :func:`~collections.namedtuple` consisting of the following fields:\n - **optim_state** - current optimizer's state.\n - **rng_key** - random number generator seed used for the iteration.\n\"\"\"\n\n\nSVIRunResult = namedtuple('SVIRunResult', ['params', 'losses'])\n\"\"\"\nA :func:`~collections.namedtuple` consisting of the following fields:\n - **params** - the optimized parameters.\n - **losses** - the losses collected at every step.\n\"\"\"\n\n\ndef _apply_loss_fn(loss_fn, rng_key, constrain_fn, model, guide,\n                   args, kwargs, static_kwargs, params):\n    return loss_fn(rng_key, constrain_fn(params), model, guide, *args, **kwargs, **static_kwargs)\n\n\nclass SVI(object):\n    \"\"\"\n    Stochastic Variational Inference given an ELBO loss objective.\n\n    **References**\n\n    1. *SVI Part I: An Introduction to Stochastic Variational Inference in Pyro*,\n       (http://pyro.ai/examples/svi_part_i.html)\n\n    **Example:**\n\n    .. doctest::\n\n        >>> from jax import random\n        >>> import jax.numpy as jnp\n        >>> import numpyro\n        >>> import numpyro.distributions as dist\n        >>> from numpyro.distributions import constraints\n        >>> from numpyro.infer import SVI, Trace_ELBO\n\n        >>> def model(data):\n        ...     f = numpyro.sample(\"latent_fairness\", dist.Beta(10, 10))\n        ...     with numpyro.plate(\"N\", data.shape[0]):\n        ...         numpyro.sample(\"obs\", dist.Bernoulli(f), obs=data)\n\n        >>> def guide(data):\n        ...     alpha_q = numpyro.param(\"alpha_q\", 15., constraint=constraints.positive)\n        ...     beta_q = numpyro.param(\"beta_q\", lambda rng_key: random.exponential(rng_key),\n        ...                            constraint=constraints.positive)\n        ...     numpyro.sample(\"latent_fairness\", dist.Beta(alpha_q, beta_q))\n\n        >>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])\n        >>> optimizer = numpyro.optim.Adam(step_size=0.0005)\n        >>> svi = SVI(model, guide, optimizer, loss=Trace_ELBO())\n        >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)\n        >>> params = svi_result.params\n        >>> inferred_mean = params[\"alpha_q\"] / (params[\"alpha_q\"] + params[\"beta_q\"])\n\n    :param model: Python callable with Pyro primitives for the model.\n    :param guide: Python callable with Pyro primitives for the guide\n        (recognition network).\n    :param optim: an instance of :class:`~numpyro.optim._NumpyroOptim`.\n    :param loss: ELBO loss, i.e. negative Evidence Lower Bound, to minimize.\n    :param static_kwargs: static arguments for the model / guide, i.e. arguments\n        that remain constant during fitting.\n    :return: tuple of `(init_fn, update_fn, evaluate)`.\n    \"\"\"\n    def __init__(self, model, guide, optim, loss, **static_kwargs):\n        self.model = model\n        self.guide = guide\n        self.loss = loss\n        self.optim = optim\n        self.static_kwargs = static_kwargs\n        self.constrain_fn = None\n\n    def init(self, rng_key, *args, **kwargs):\n        \"\"\"\n        Gets the initial SVI state.\n\n        :param jax.random.PRNGKey rng_key: random number generator seed.\n        :param args: arguments to the model / guide (these can possibly vary during\n            the course of fitting).\n        :param kwargs: keyword arguments to the model / guide (these can possibly vary\n            during the course of fitting).\n        :return: the initial :data:`SVIState`\n        \"\"\"\n        rng_key, model_seed, guide_seed = random.split(rng_key, 3)\n        model_init = seed(self.model, model_seed)\n        guide_init = seed(self.guide, guide_seed)\n        guide_trace = trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)\n        model_trace = trace(replay(model_init, guide_trace)).get_trace(*args, **kwargs, **self.static_kwargs)\n        params = {}\n        inv_transforms = {}\n        # NB: params in model_trace will be overwritten by params in guide_trace\n        for site in list(model_trace.values()) + list(guide_trace.values()):\n            if site['type'] == 'param':\n                constraint = site['kwargs'].pop('constraint', constraints.real)\n                transform = biject_to(constraint)\n                inv_transforms[site['name']] = transform\n                params[site['name']] = transform.inv(site['value'])\n\n        self.constrain_fn = partial(transform_fn, inv_transforms)\n        return SVIState(self.optim.init(params), rng_key)\n\n    def get_params(self, svi_state):\n        \"\"\"\n        Gets values at `param` sites of the `model` and `guide`.\n\n        :param svi_state: current state of SVI.\n        :return: the corresponding parameters\n        \"\"\"\n        params = self.constrain_fn(self.optim.get_params(svi_state.optim_state))\n        return params\n\n    def update(self, svi_state, *args, **kwargs):\n        \"\"\"\n        Take a single step of SVI (possibly on a batch / minibatch of data),\n        using the optimizer.\n\n        :param svi_state: current state of SVI.\n        :param args: arguments to the model / guide (these can possibly vary during\n            the course of fitting).\n        :param kwargs: keyword arguments to the model / guide (these can possibly vary\n            during the course of fitting).\n        :return: tuple of `(svi_state, loss)`.\n        \"\"\"\n        rng_key, rng_key_step = random.split(svi_state.rng_key)\n        loss_fn = partial(_apply_loss_fn, self.loss.loss, rng_key_step, self.constrain_fn, self.model,\n                          self.guide, args, kwargs, self.static_kwargs)\n        loss_val, optim_state = self.optim.eval_and_update(loss_fn, svi_state.optim_state)\n        return SVIState(optim_state, rng_key), loss_val\n\n    def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):\n        \"\"\"\n        (EXPERIMENTAL INTERFACE) Run SVI with `num_steps` iterations, then return\n        the optimized parameters and the stacked losses at every step. If `num_steps`\n        is large, setting `progress_bar=False` can make the run faster.\n\n        .. note:: For a complex training process (e.g. the one requires early stopping,\n            epoch training, varying args/kwargs,...), we recommend to use the more\n            flexible methods :meth:`init`, :meth:`update`, :meth:`evaluate` to\n            customize your training procedure.\n\n        :param jax.random.PRNGKey rng_key: random number generator seed.\n        :param int num_steps: the number of optimization steps.\n        :param args: arguments to the model / guide\n        :param bool progress_bar: Whether to enable progress bar updates. Defaults to\n            ``True``.\n        :param kwargs: keyword arguments to the model / guide\n        :return: a namedtuple with fields `params` and `losses` where `params`\n            holds the optimized values at :class:`numpyro.param` sites,\n            and `losses` is the collected loss during the process.\n        :rtype: SVIRunResult\n        \"\"\"\n        def body_fn(svi_state, carry):\n            svi_state, loss = self.update(svi_state, *args, **kwargs)\n            return svi_state, loss\n\n        svi_state = self.init(rng_key, *args, **kwargs)\n        if progress_bar:\n            losses = []\n            with tqdm.trange(1, num_steps + 1) as t:\n                batch = num_steps // 20\n                for i in t:\n                    svi_state, loss = jit(body_fn)(svi_state, None)\n                    if i % batch == 0:\n                        avg_loss = sum(losses[i-batch:]) / batch\n                        t.set_postfix_str(\"init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}\"\n                                          .format(losses[0], i - batch + 1, i, avg_loss),\n                                          refresh=False)\n                    losses.append(loss)\n            losses = jnp.stack(losses)\n        else:\n            svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)\n\n        return SVIRunResult(self.get_params(svi_state), losses)\n\n    def evaluate(self, svi_state, *args, **kwargs):\n        \"\"\"\n        Take a single step of SVI (possibly on a batch / minibatch of data).\n\n        :param svi_state: current state of SVI.\n        :param args: arguments to the model / guide (these can possibly vary during\n            the course of fitting).\n        :param kwargs: keyword arguments to the model / guide.\n        :return: evaluate ELBO loss given the current parameter values\n            (held within `svi_state.optim_state`).\n        \"\"\"\n        # we split to have the same seed as `update_fn` given an svi_state\n        _, rng_key_eval = random.split(svi_state.rng_key)\n        params = self.get_params(svi_state)\n        return self.loss.loss(rng_key_eval, params, self.model, self.guide,\n                              *args, **kwargs, **self.static_kwargs)\n", "path": "numpyro/infer/svi.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nfrom functools import namedtuple, partial\n\nimport tqdm\n\nfrom jax import jit, lax, random\nimport jax.numpy as jnp\n\nfrom numpyro.distributions import constraints\nfrom numpyro.distributions.transforms import biject_to\nfrom numpyro.handlers import replay, seed, trace\nfrom numpyro.infer.util import transform_fn\n\nSVIState = namedtuple('SVIState', ['optim_state', 'rng_key'])\n\"\"\"\nA :func:`~collections.namedtuple` consisting of the following fields:\n - **optim_state** - current optimizer's state.\n - **rng_key** - random number generator seed used for the iteration.\n\"\"\"\n\n\nSVIRunResult = namedtuple('SVIRunResult', ['params', 'losses'])\n\"\"\"\nA :func:`~collections.namedtuple` consisting of the following fields:\n - **params** - the optimized parameters.\n - **losses** - the losses collected at every step.\n\"\"\"\n\n\ndef _apply_loss_fn(loss_fn, rng_key, constrain_fn, model, guide,\n                   args, kwargs, static_kwargs, params):\n    return loss_fn(rng_key, constrain_fn(params), model, guide, *args, **kwargs, **static_kwargs)\n\n\nclass SVI(object):\n    \"\"\"\n    Stochastic Variational Inference given an ELBO loss objective.\n\n    **References**\n\n    1. *SVI Part I: An Introduction to Stochastic Variational Inference in Pyro*,\n       (http://pyro.ai/examples/svi_part_i.html)\n\n    **Example:**\n\n    .. doctest::\n\n        >>> from jax import random\n        >>> import jax.numpy as jnp\n        >>> import numpyro\n        >>> import numpyro.distributions as dist\n        >>> from numpyro.distributions import constraints\n        >>> from numpyro.infer import SVI, Trace_ELBO\n\n        >>> def model(data):\n        ...     f = numpyro.sample(\"latent_fairness\", dist.Beta(10, 10))\n        ...     with numpyro.plate(\"N\", data.shape[0]):\n        ...         numpyro.sample(\"obs\", dist.Bernoulli(f), obs=data)\n\n        >>> def guide(data):\n        ...     alpha_q = numpyro.param(\"alpha_q\", 15., constraint=constraints.positive)\n        ...     beta_q = numpyro.param(\"beta_q\", lambda rng_key: random.exponential(rng_key),\n        ...                            constraint=constraints.positive)\n        ...     numpyro.sample(\"latent_fairness\", dist.Beta(alpha_q, beta_q))\n\n        >>> data = jnp.concatenate([jnp.ones(6), jnp.zeros(4)])\n        >>> optimizer = numpyro.optim.Adam(step_size=0.0005)\n        >>> svi = SVI(model, guide, optimizer, loss=Trace_ELBO())\n        >>> svi_result = svi.run(random.PRNGKey(0), 2000, data)\n        >>> params = svi_result.params\n        >>> inferred_mean = params[\"alpha_q\"] / (params[\"alpha_q\"] + params[\"beta_q\"])\n\n    :param model: Python callable with Pyro primitives for the model.\n    :param guide: Python callable with Pyro primitives for the guide\n        (recognition network).\n    :param optim: an instance of :class:`~numpyro.optim._NumpyroOptim`.\n    :param loss: ELBO loss, i.e. negative Evidence Lower Bound, to minimize.\n    :param static_kwargs: static arguments for the model / guide, i.e. arguments\n        that remain constant during fitting.\n    :return: tuple of `(init_fn, update_fn, evaluate)`.\n    \"\"\"\n    def __init__(self, model, guide, optim, loss, **static_kwargs):\n        self.model = model\n        self.guide = guide\n        self.loss = loss\n        self.optim = optim\n        self.static_kwargs = static_kwargs\n        self.constrain_fn = None\n\n    def init(self, rng_key, *args, **kwargs):\n        \"\"\"\n        Gets the initial SVI state.\n\n        :param jax.random.PRNGKey rng_key: random number generator seed.\n        :param args: arguments to the model / guide (these can possibly vary during\n            the course of fitting).\n        :param kwargs: keyword arguments to the model / guide (these can possibly vary\n            during the course of fitting).\n        :return: the initial :data:`SVIState`\n        \"\"\"\n        rng_key, model_seed, guide_seed = random.split(rng_key, 3)\n        model_init = seed(self.model, model_seed)\n        guide_init = seed(self.guide, guide_seed)\n        guide_trace = trace(guide_init).get_trace(*args, **kwargs, **self.static_kwargs)\n        model_trace = trace(replay(model_init, guide_trace)).get_trace(*args, **kwargs, **self.static_kwargs)\n        params = {}\n        inv_transforms = {}\n        # NB: params in model_trace will be overwritten by params in guide_trace\n        for site in list(model_trace.values()) + list(guide_trace.values()):\n            if site['type'] == 'param':\n                constraint = site['kwargs'].pop('constraint', constraints.real)\n                transform = biject_to(constraint)\n                inv_transforms[site['name']] = transform\n                params[site['name']] = transform.inv(site['value'])\n\n        self.constrain_fn = partial(transform_fn, inv_transforms)\n        return SVIState(self.optim.init(params), rng_key)\n\n    def get_params(self, svi_state):\n        \"\"\"\n        Gets values at `param` sites of the `model` and `guide`.\n\n        :param svi_state: current state of SVI.\n        :return: the corresponding parameters\n        \"\"\"\n        params = self.constrain_fn(self.optim.get_params(svi_state.optim_state))\n        return params\n\n    def update(self, svi_state, *args, **kwargs):\n        \"\"\"\n        Take a single step of SVI (possibly on a batch / minibatch of data),\n        using the optimizer.\n\n        :param svi_state: current state of SVI.\n        :param args: arguments to the model / guide (these can possibly vary during\n            the course of fitting).\n        :param kwargs: keyword arguments to the model / guide (these can possibly vary\n            during the course of fitting).\n        :return: tuple of `(svi_state, loss)`.\n        \"\"\"\n        rng_key, rng_key_step = random.split(svi_state.rng_key)\n        loss_fn = partial(_apply_loss_fn, self.loss.loss, rng_key_step, self.constrain_fn, self.model,\n                          self.guide, args, kwargs, self.static_kwargs)\n        loss_val, optim_state = self.optim.eval_and_update(loss_fn, svi_state.optim_state)\n        return SVIState(optim_state, rng_key), loss_val\n\n    def run(self, rng_key, num_steps, *args, progress_bar=True, **kwargs):\n        \"\"\"\n        (EXPERIMENTAL INTERFACE) Run SVI with `num_steps` iterations, then return\n        the optimized parameters and the stacked losses at every step. If `num_steps`\n        is large, setting `progress_bar=False` can make the run faster.\n\n        .. note:: For a complex training process (e.g. the one requires early stopping,\n            epoch training, varying args/kwargs,...), we recommend to use the more\n            flexible methods :meth:`init`, :meth:`update`, :meth:`evaluate` to\n            customize your training procedure.\n\n        :param jax.random.PRNGKey rng_key: random number generator seed.\n        :param int num_steps: the number of optimization steps.\n        :param args: arguments to the model / guide\n        :param bool progress_bar: Whether to enable progress bar updates. Defaults to\n            ``True``.\n        :param kwargs: keyword arguments to the model / guide\n        :return: a namedtuple with fields `params` and `losses` where `params`\n            holds the optimized values at :class:`numpyro.param` sites,\n            and `losses` is the collected loss during the process.\n        :rtype: SVIRunResult\n        \"\"\"\n        def body_fn(svi_state, carry):\n            svi_state, loss = self.update(svi_state, *args, **kwargs)\n            return svi_state, loss\n\n        svi_state = self.init(rng_key, *args, **kwargs)\n        if progress_bar:\n            losses = []\n            with tqdm.trange(1, num_steps + 1) as t:\n                batch = max(num_steps // 20, 1)\n                for i in t:\n                    svi_state, loss = jit(body_fn)(svi_state, None)\n                    losses.append(loss)\n                    if i % batch == 0:\n                        avg_loss = sum(losses[i-batch:]) / batch\n                        t.set_postfix_str(\"init loss: {:.4f}, avg. loss [{}-{}]: {:.4f}\"\n                                          .format(losses[0], i - batch + 1, i, avg_loss),\n                                          refresh=False)\n            losses = jnp.stack(losses)\n        else:\n            svi_state, losses = lax.scan(body_fn, svi_state, None, length=num_steps)\n\n        return SVIRunResult(self.get_params(svi_state), losses)\n\n    def evaluate(self, svi_state, *args, **kwargs):\n        \"\"\"\n        Take a single step of SVI (possibly on a batch / minibatch of data).\n\n        :param svi_state: current state of SVI.\n        :param args: arguments to the model / guide (these can possibly vary during\n            the course of fitting).\n        :param kwargs: keyword arguments to the model / guide.\n        :return: evaluate ELBO loss given the current parameter values\n            (held within `svi_state.optim_state`).\n        \"\"\"\n        # we split to have the same seed as `update_fn` given an svi_state\n        _, rng_key_eval = random.split(svi_state.rng_key)\n        params = self.get_params(svi_state)\n        return self.loss.loss(rng_key_eval, params, self.model, self.guide,\n                              *args, **kwargs, **self.static_kwargs)\n", "path": "numpyro/infer/svi.py"}]}