prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Update docstring of Neal's funnel example
  We have updated [funnel](https://github.com/pyro-ppl/numpyro/blob/master/examples/funnel.py) example to use `reparam` handler, but the docstring is not updated yet.
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `examples/funnel.py`
  Content:
  ```
  1 # Copyright Contributors to the Pyro project.
  2 # SPDX-License-Identifier: Apache-2.0
  3 
  4 """
  5 Example: Neal's Funnel
  6 ======================
  7 
  8 This example, which is adapted from [1], illustrates how to leverage non-centered
  9 parameterization using the class :class:`numpyro.distributions.TransformedDistribution`.
  10 We will examine the difference between two types of parameterizations on the
  11 10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at
  12 the neck of the funnel if centered parameterization is used. On the contrary,
  13 the problem can be solved by using non-centered parameterization.
  14 
  15 Using non-centered parameterization through TransformedDistribution in NumPyro
  16 has the same effect as the automatic reparameterisation technique introduced in
  17 [2]. However, in [2], users need to implement a (non-trivial) reparameterization
  18 rule for each type of transform. Instead, in NumPyro the only requirement to let
  19 inference algorithms know to do reparameterization automatically is to declare
  20 the random variable as a transformed distribution.
  21 
  22 **References:**
  23 
  24     1. *Stan User's Guide*, https://mc-stan.org/docs/2_19/stan-users-guide/reparameterization-section.html
  25     2. Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019), "Automatic
  26        Reparameterisation of Probabilistic Programs", (https://arxiv.org/abs/1906.03028)
  27 """
  28 
  29 import argparse
  30 import os
  31 
  32 import matplotlib.pyplot as plt
  33 
  34 from jax import random
  35 import jax.numpy as jnp
  36 
  37 import numpyro
  38 import numpyro.distributions as dist
  39 from numpyro.infer import MCMC, NUTS, Predictive
  40 from numpyro.infer.reparam import LocScaleReparam
  41 
  42 
  43 def model(dim=10):
  44     y = numpyro.sample('y', dist.Normal(0, 3))
  45     numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))
  46 
  47 
  48 def reparam_model(dim=10):
  49     y = numpyro.sample('y', dist.Normal(0, 3))
  50     with numpyro.handlers.reparam(config={'x': LocScaleReparam(0)}):
  51         numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))
  52 
  53 
  54 def run_inference(model, args, rng_key):
  55     kernel = NUTS(model)
  56     mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
  57                 progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
  58     mcmc.run(rng_key)
  59     mcmc.print_summary()
  60     return mcmc.get_samples()
  61 
  62 
  63 def main(args):
  64     rng_key = random.PRNGKey(0)
  65 
  66     # do inference with centered parameterization
  67     print("============================= Centered Parameterization ==============================")
  68     samples = run_inference(model, args, rng_key)
  69 
  70     # do inference with non-centered parameterization
  71     print("\n=========================== Non-centered Parameterization ============================")
  72     reparam_samples = run_inference(reparam_model, args, rng_key)
  73     # collect deterministic sites
  74     reparam_samples = Predictive(reparam_model, reparam_samples, return_sites=['x', 'y'])(
  75         random.PRNGKey(1))
  76 
  77     # make plots
  78     fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 8))
  79 
  80     ax1.plot(samples['x'][:, 0], samples['y'], "go", alpha=0.3)
  81     ax1.set(xlim=(-20, 20), ylim=(-9, 9), ylabel='y',
  82             title='Funnel samples with centered parameterization')
  83 
  84     ax2.plot(reparam_samples['x'][:, 0], reparam_samples['y'], "go", alpha=0.3)
  85     ax2.set(xlim=(-20, 20), ylim=(-9, 9), xlabel='x[0]', ylabel='y',
  86             title='Funnel samples with non-centered parameterization')
  87 
  88     plt.savefig('funnel_plot.pdf')
  89     plt.tight_layout()
  90 
  91 
  92 if __name__ == "__main__":
  93     assert numpyro.__version__.startswith('0.4.1')
  94     parser = argparse.ArgumentParser(description="Non-centered reparameterization example")
  95     parser.add_argument("-n", "--num-samples", nargs="?", default=1000, type=int)
  96     parser.add_argument("--num-warmup", nargs='?', default=1000, type=int)
  97     parser.add_argument("--num-chains", nargs='?', default=1, type=int)
  98     parser.add_argument("--device", default='cpu', type=str, help='use "cpu" or "gpu".')
  99     args = parser.parse_args()
  100 
  101     numpyro.set_platform(args.device)
  102     numpyro.set_host_device_count(args.num_chains)
  103 
  104     main(args)
  105
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/examples/funnel.py b/examples/funnel.py
  --- a/examples/funnel.py
  +++ b/examples/funnel.py
  @@ -6,18 +6,15 @@
   ======================
   
   This example, which is adapted from [1], illustrates how to leverage non-centered
  -parameterization using the class :class:`numpyro.distributions.TransformedDistribution`.
  +parameterization using the :class:`~numpyro.handlers.reparam` handler.
   We will examine the difference between two types of parameterizations on the
   10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at
   the neck of the funnel if centered parameterization is used. On the contrary,
   the problem can be solved by using non-centered parameterization.
   
  -Using non-centered parameterization through TransformedDistribution in NumPyro
  -has the same effect as the automatic reparameterisation technique introduced in
  -[2]. However, in [2], users need to implement a (non-trivial) reparameterization
  -rule for each type of transform. Instead, in NumPyro the only requirement to let
  -inference algorithms know to do reparameterization automatically is to declare
  -the random variable as a transformed distribution.
  +Using non-centered parameterization through :class:`~numpyro.infer.reparam.LocScaleReparam`
  +or :class:`~numpyro.infer.reparam.TransformReparam` in NumPyro has the same effect as
  +the automatic reparameterisation technique introduced in [2].
   
   **References:**
   
  @@ -36,6 +33,7 @@
   
   import numpyro
   import numpyro.distributions as dist
  +from numpyro.handlers import reparam
   from numpyro.infer import MCMC, NUTS, Predictive
   from numpyro.infer.reparam import LocScaleReparam
   
  @@ -45,10 +43,7 @@
       numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))
   
   
  -def reparam_model(dim=10):
  -    y = numpyro.sample('y', dist.Normal(0, 3))
  -    with numpyro.handlers.reparam(config={'x': LocScaleReparam(0)}):
  -        numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))
  +reparam_model = reparam(model, config={'x': LocScaleReparam(0)})
   
   
   def run_inference(model, args, rng_key):
  @@ -56,7 +51,7 @@
       mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,
                   progress_bar=False if "NUMPYRO_SPHINXBUILD" in os.environ else True)
       mcmc.run(rng_key)
  -    mcmc.print_summary()
  +    mcmc.print_summary(exclude_deterministic=False)
       return mcmc.get_samples()
  
verification_info: {"golden_diff": "diff --git a/examples/funnel.py b/examples/funnel.py\n--- a/examples/funnel.py\n+++ b/examples/funnel.py\n@@ -6,18 +6,15 @@\n ======================\n \n This example, which is adapted from [1], illustrates how to leverage non-centered\n-parameterization using the class :class:`numpyro.distributions.TransformedDistribution`.\n+parameterization using the :class:`~numpyro.handlers.reparam` handler.\n We will examine the difference between two types of parameterizations on the\n 10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at\n the neck of the funnel if centered parameterization is used. On the contrary,\n the problem can be solved by using non-centered parameterization.\n \n-Using non-centered parameterization through TransformedDistribution in NumPyro\n-has the same effect as the automatic reparameterisation technique introduced in\n-[2]. However, in [2], users need to implement a (non-trivial) reparameterization\n-rule for each type of transform. Instead, in NumPyro the only requirement to let\n-inference algorithms know to do reparameterization automatically is to declare\n-the random variable as a transformed distribution.\n+Using non-centered parameterization through :class:`~numpyro.infer.reparam.LocScaleReparam`\n+or :class:`~numpyro.infer.reparam.TransformReparam` in NumPyro has the same effect as\n+the automatic reparameterisation technique introduced in [2].\n \n **References:**\n \n@@ -36,6 +33,7 @@\n \n import numpyro\n import numpyro.distributions as dist\n+from numpyro.handlers import reparam\n from numpyro.infer import MCMC, NUTS, Predictive\n from numpyro.infer.reparam import LocScaleReparam\n \n@@ -45,10 +43,7 @@\n     numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))\n \n \n-def reparam_model(dim=10):\n-    y = numpyro.sample('y', dist.Normal(0, 3))\n-    with numpyro.handlers.reparam(config={'x': LocScaleReparam(0)}):\n-        numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))\n+reparam_model = reparam(model, config={'x': LocScaleReparam(0)})\n \n \n def run_inference(model, args, rng_key):\n@@ -56,7 +51,7 @@\n     mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,\n                 progress_bar=False if \"NUMPYRO_SPHINXBUILD\" in os.environ else True)\n     mcmc.run(rng_key)\n-    mcmc.print_summary()\n+    mcmc.print_summary(exclude_deterministic=False)\n     return mcmc.get_samples()\n", "issue": "Update docstring of Neal's funnel example\nWe have updated [funnel](https://github.com/pyro-ppl/numpyro/blob/master/examples/funnel.py) example to use `reparam` handler, but the docstring is not updated yet.\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\n\"\"\"\nExample: Neal's Funnel\n======================\n\nThis example, which is adapted from [1], illustrates how to leverage non-centered\nparameterization using the class :class:`numpyro.distributions.TransformedDistribution`.\nWe will examine the difference between two types of parameterizations on the\n10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at\nthe neck of the funnel if centered parameterization is used. On the contrary,\nthe problem can be solved by using non-centered parameterization.\n\nUsing non-centered parameterization through TransformedDistribution in NumPyro\nhas the same effect as the automatic reparameterisation technique introduced in\n[2]. However, in [2], users need to implement a (non-trivial) reparameterization\nrule for each type of transform. Instead, in NumPyro the only requirement to let\ninference algorithms know to do reparameterization automatically is to declare\nthe random variable as a transformed distribution.\n\n**References:**\n\n    1. *Stan User's Guide*, https://mc-stan.org/docs/2_19/stan-users-guide/reparameterization-section.html\n    2. Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019), \"Automatic\n       Reparameterisation of Probabilistic Programs\", (https://arxiv.org/abs/1906.03028)\n\"\"\"\n\nimport argparse\nimport os\n\nimport matplotlib.pyplot as plt\n\nfrom jax import random\nimport jax.numpy as jnp\n\nimport numpyro\nimport numpyro.distributions as dist\nfrom numpyro.infer import MCMC, NUTS, Predictive\nfrom numpyro.infer.reparam import LocScaleReparam\n\n\ndef model(dim=10):\n    y = numpyro.sample('y', dist.Normal(0, 3))\n    numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))\n\n\ndef reparam_model(dim=10):\n    y = numpyro.sample('y', dist.Normal(0, 3))\n    with numpyro.handlers.reparam(config={'x': LocScaleReparam(0)}):\n        numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))\n\n\ndef run_inference(model, args, rng_key):\n    kernel = NUTS(model)\n    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,\n                progress_bar=False if \"NUMPYRO_SPHINXBUILD\" in os.environ else True)\n    mcmc.run(rng_key)\n    mcmc.print_summary()\n    return mcmc.get_samples()\n\n\ndef main(args):\n    rng_key = random.PRNGKey(0)\n\n    # do inference with centered parameterization\n    print(\"============================= Centered Parameterization ==============================\")\n    samples = run_inference(model, args, rng_key)\n\n    # do inference with non-centered parameterization\n    print(\"\\n=========================== Non-centered Parameterization ============================\")\n    reparam_samples = run_inference(reparam_model, args, rng_key)\n    # collect deterministic sites\n    reparam_samples = Predictive(reparam_model, reparam_samples, return_sites=['x', 'y'])(\n        random.PRNGKey(1))\n\n    # make plots\n    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 8))\n\n    ax1.plot(samples['x'][:, 0], samples['y'], \"go\", alpha=0.3)\n    ax1.set(xlim=(-20, 20), ylim=(-9, 9), ylabel='y',\n            title='Funnel samples with centered parameterization')\n\n    ax2.plot(reparam_samples['x'][:, 0], reparam_samples['y'], \"go\", alpha=0.3)\n    ax2.set(xlim=(-20, 20), ylim=(-9, 9), xlabel='x[0]', ylabel='y',\n            title='Funnel samples with non-centered parameterization')\n\n    plt.savefig('funnel_plot.pdf')\n    plt.tight_layout()\n\n\nif __name__ == \"__main__\":\n    assert numpyro.__version__.startswith('0.4.1')\n    parser = argparse.ArgumentParser(description=\"Non-centered reparameterization example\")\n    parser.add_argument(\"-n\", \"--num-samples\", nargs=\"?\", default=1000, type=int)\n    parser.add_argument(\"--num-warmup\", nargs='?', default=1000, type=int)\n    parser.add_argument(\"--num-chains\", nargs='?', default=1, type=int)\n    parser.add_argument(\"--device\", default='cpu', type=str, help='use \"cpu\" or \"gpu\".')\n    args = parser.parse_args()\n\n    numpyro.set_platform(args.device)\n    numpyro.set_host_device_count(args.num_chains)\n\n    main(args)\n", "path": "examples/funnel.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\n\"\"\"\nExample: Neal's Funnel\n======================\n\nThis example, which is adapted from [1], illustrates how to leverage non-centered\nparameterization using the :class:`~numpyro.handlers.reparam` handler.\nWe will examine the difference between two types of parameterizations on the\n10-dimensional Neal's funnel distribution. As we will see, HMC gets trouble at\nthe neck of the funnel if centered parameterization is used. On the contrary,\nthe problem can be solved by using non-centered parameterization.\n\nUsing non-centered parameterization through :class:`~numpyro.infer.reparam.LocScaleReparam`\nor :class:`~numpyro.infer.reparam.TransformReparam` in NumPyro has the same effect as\nthe automatic reparameterisation technique introduced in [2].\n\n**References:**\n\n    1. *Stan User's Guide*, https://mc-stan.org/docs/2_19/stan-users-guide/reparameterization-section.html\n    2. Maria I. Gorinova, Dave Moore, Matthew D. Hoffman (2019), \"Automatic\n       Reparameterisation of Probabilistic Programs\", (https://arxiv.org/abs/1906.03028)\n\"\"\"\n\nimport argparse\nimport os\n\nimport matplotlib.pyplot as plt\n\nfrom jax import random\nimport jax.numpy as jnp\n\nimport numpyro\nimport numpyro.distributions as dist\nfrom numpyro.handlers import reparam\nfrom numpyro.infer import MCMC, NUTS, Predictive\nfrom numpyro.infer.reparam import LocScaleReparam\n\n\ndef model(dim=10):\n    y = numpyro.sample('y', dist.Normal(0, 3))\n    numpyro.sample('x', dist.Normal(jnp.zeros(dim - 1), jnp.exp(y / 2)))\n\n\nreparam_model = reparam(model, config={'x': LocScaleReparam(0)})\n\n\ndef run_inference(model, args, rng_key):\n    kernel = NUTS(model)\n    mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains,\n                progress_bar=False if \"NUMPYRO_SPHINXBUILD\" in os.environ else True)\n    mcmc.run(rng_key)\n    mcmc.print_summary(exclude_deterministic=False)\n    return mcmc.get_samples()\n\n\ndef main(args):\n    rng_key = random.PRNGKey(0)\n\n    # do inference with centered parameterization\n    print(\"============================= Centered Parameterization ==============================\")\n    samples = run_inference(model, args, rng_key)\n\n    # do inference with non-centered parameterization\n    print(\"\\n=========================== Non-centered Parameterization ============================\")\n    reparam_samples = run_inference(reparam_model, args, rng_key)\n    # collect deterministic sites\n    reparam_samples = Predictive(reparam_model, reparam_samples, return_sites=['x', 'y'])(\n        random.PRNGKey(1))\n\n    # make plots\n    fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, figsize=(8, 8))\n\n    ax1.plot(samples['x'][:, 0], samples['y'], \"go\", alpha=0.3)\n    ax1.set(xlim=(-20, 20), ylim=(-9, 9), ylabel='y',\n            title='Funnel samples with centered parameterization')\n\n    ax2.plot(reparam_samples['x'][:, 0], reparam_samples['y'], \"go\", alpha=0.3)\n    ax2.set(xlim=(-20, 20), ylim=(-9, 9), xlabel='x[0]', ylabel='y',\n            title='Funnel samples with non-centered parameterization')\n\n    plt.savefig('funnel_plot.pdf')\n    plt.tight_layout()\n\n\nif __name__ == \"__main__\":\n    assert numpyro.__version__.startswith('0.4.1')\n    parser = argparse.ArgumentParser(description=\"Non-centered reparameterization example\")\n    parser.add_argument(\"-n\", \"--num-samples\", nargs=\"?\", default=1000, type=int)\n    parser.add_argument(\"--num-warmup\", nargs='?', default=1000, type=int)\n    parser.add_argument(\"--num-chains\", nargs='?', default=1, type=int)\n    parser.add_argument(\"--device\", default='cpu', type=str, help='use \"cpu\" or \"gpu\".')\n    args = parser.parse_args()\n\n    numpyro.set_platform(args.device)\n    numpyro.set_host_device_count(args.num_chains)\n\n    main(args)\n", "path": "examples/funnel.py"}]}