We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
compat module: compat.infer.MCMC broken
Calling
```python
compat.infer.MCMC(kernel, num_warmup=100, num_samples=1000)
```
throws an error since no kwargs are handed to `mcmc.MCMC`. Also, shouldnt the correct `compat` call be
```python
compat.infer.MCMC(kernel, warmup_steps=100, num_samples=1000)
```
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/compat/infer.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 import math
5 
6 from jax import jit
7 
8 import numpyro
9 from numpyro.compat.pyro import get_param_store
10 from numpyro.infer import elbo, hmc, mcmc, svi
11 
12 
13 class HMC(hmc.HMC):
14     def __init__(
15         self,
16         model=None,
17         potential_fn=None,
18         step_size=1,
19         adapt_step_size=True,
20         adapt_mass_matrix=True,
21         full_mass=False,
22         use_multinomial_sampling=True,
23         transforms=None,
24         max_plate_nesting=None,
25         jit_compile=False,
26         jit_options=None,
27         ignore_jit_warnings=False,
28         trajectory_length=2 * math.pi,
29         target_accept_prob=0.8,
30     ):
31         super(HMC, self).__init__(
32             model=model,
33             potential_fn=potential_fn,
34             step_size=step_size,
35             adapt_step_size=adapt_step_size,
36             adapt_mass_matrix=adapt_mass_matrix,
37             dense_mass=full_mass,
38             target_accept_prob=target_accept_prob,
39             trajectory_length=trajectory_length,
40         )
41 
42 
43 class NUTS(hmc.NUTS):
44     def __init__(
45         self,
46         model=None,
47         potential_fn=None,
48         step_size=1,
49         adapt_step_size=True,
50         adapt_mass_matrix=True,
51         full_mass=False,
52         use_multinomial_sampling=True,
53         transforms=None,
54         max_plate_nesting=None,
55         jit_compile=False,
56         jit_options=None,
57         ignore_jit_warnings=False,
58         trajectory_length=2 * math.pi,
59         target_accept_prob=0.8,
60         max_tree_depth=10,
61     ):
62         if potential_fn is not None:
63             raise ValueError(
64                 "Only `model` argument is supported in generic module;"
65                 " `potential_fn` is not supported."
66             )
67         super(NUTS, self).__init__(
68             model=model,
69             potential_fn=potential_fn,
70             step_size=step_size,
71             adapt_step_size=adapt_step_size,
72             adapt_mass_matrix=adapt_mass_matrix,
73             dense_mass=full_mass,
74             target_accept_prob=target_accept_prob,
75             trajectory_length=trajectory_length,
76             max_tree_depth=max_tree_depth,
77         )
78 
79 
80 class MCMC(object):
81     def __init__(
82         self,
83         kernel,
84         num_samples,
85         num_warmup=None,
86         initial_params=None,
87         num_chains=1,
88         hook_fn=None,
89         mp_context=None,
90         disable_progbar=False,
91         disable_validation=True,
92         transforms=None,
93     ):
94         if num_warmup is None:
95             num_warmup = num_samples
96         self._initial_params = initial_params
97         self._mcmc = mcmc.MCMC(
98             kernel,
99             num_warmup,
100             num_samples,
101             num_chains=num_chains,
102             progress_bar=(not disable_progbar),
103         )
104 
105     def run(self, *args, rng_key=None, **kwargs):
106         if rng_key is None:
107             rng_key = numpyro.prng_key()
108         self._mcmc.run(rng_key, *args, init_params=self._initial_params, **kwargs)
109 
110     def get_samples(self, num_samples=None, group_by_chain=False):
111         if num_samples is not None:
112             raise ValueError("`num_samples` arg unsupported in NumPyro.")
113         return self._mcmc.get_samples(group_by_chain=group_by_chain)
114 
115     def summary(self, prob=0.9):
116         self._mcmc.print_summary()
117 
118 
119 class SVI(svi.SVI):
120     def __init__(
121         self,
122         model,
123         guide,
124         optim,
125         loss,
126         loss_and_grads=None,
127         num_samples=10,
128         num_steps=0,
129         **kwargs
130     ):
131         super(SVI, self).__init__(model=model, guide=guide, optim=optim, loss=loss)
132         self.svi_state = None
133 
134     def evaluate_loss(self, *args, **kwargs):
135         return self.evaluate(self.svi_state, *args, **kwargs)
136 
137     def step(self, *args, rng_key=None, **kwargs):
138         if self.svi_state is None:
139             if rng_key is None:
140                 rng_key = numpyro.prng_key()
141             self.svi_state = self.init(rng_key, *args, **kwargs)
142         try:
143             self.svi_state, loss = jit(self.update)(self.svi_state, *args, **kwargs)
144         except TypeError as e:
145             if "not a valid JAX type" in str(e):
146                 raise TypeError(
147                     "NumPyro backend requires args, kwargs to be arrays or tuples, "
148                     "dicts of arrays."
149                 ) from e
150             else:
151                 raise e
152         params = jit(super(SVI, self).get_params)(self.svi_state)
153         get_param_store().update(params)
154         return loss
155 
156     def get_params(self):
157         return super(SVI, self).get_params(self.svi_state)
158 
159 
160 class Trace_ELBO(elbo.Trace_ELBO):
161     def __init__(
162         self,
163         num_particles=1,
164         max_plate_nesting=float("inf"),
165         max_iarange_nesting=None,  # DEPRECATED
166         vectorize_particles=False,
167         strict_enumeration_warning=True,
168         ignore_jit_warnings=False,
169         jit_options=None,
170         retain_graph=None,
171         tail_adaptive_beta=-1.0,
172     ):
173         super(Trace_ELBO, self).__init__(num_particles=num_particles)
174 
175 
176 # JIT is enabled by default
177 JitTrace_ELBO = Trace_ELBO
178
```
--- 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/compat/infer.py b/numpyro/compat/infer.py
--- a/numpyro/compat/infer.py
+++ b/numpyro/compat/infer.py
@@ -82,7 +82,7 @@
         self,
         kernel,
         num_samples,
-        num_warmup=None,
+        warmup_steps=None,
         initial_params=None,
         num_chains=1,
         hook_fn=None,
@@ -91,13 +91,13 @@
         disable_validation=True,
         transforms=None,
     ):
-        if num_warmup is None:
-            num_warmup = num_samples
+        if warmup_steps is None:
+            warmup_steps = num_samples
         self._initial_params = initial_params
         self._mcmc = mcmc.MCMC(
             kernel,
-            num_warmup,
-            num_samples,
+            num_warmup=warmup_steps,
+            num_samples=num_samples,
             num_chains=num_chains,
             progress_bar=(not disable_progbar),
         )

{"golden_diff": "diff --git a/numpyro/compat/infer.py b/numpyro/compat/infer.py\n--- a/numpyro/compat/infer.py\n+++ b/numpyro/compat/infer.py\n@@ -82,7 +82,7 @@\n         self,\n         kernel,\n         num_samples,\n-        num_warmup=None,\n+        warmup_steps=None,\n         initial_params=None,\n         num_chains=1,\n         hook_fn=None,\n@@ -91,13 +91,13 @@\n         disable_validation=True,\n         transforms=None,\n     ):\n-        if num_warmup is None:\n-            num_warmup = num_samples\n+        if warmup_steps is None:\n+            warmup_steps = num_samples\n         self._initial_params = initial_params\n         self._mcmc = mcmc.MCMC(\n             kernel,\n-            num_warmup,\n-            num_samples,\n+            num_warmup=warmup_steps,\n+            num_samples=num_samples,\n             num_chains=num_chains,\n             progress_bar=(not disable_progbar),\n         )\n", "issue": "compat module: compat.infer.MCMC broken\nCalling\r\n```python\r\ncompat.infer.MCMC(kernel, num_warmup=100, num_samples=1000)\r\n```\r\nthrows an error since no kwargs are handed to `mcmc.MCMC`. Also, shouldnt the correct `compat` call be\r\n```python\r\ncompat.infer.MCMC(kernel, warmup_steps=100, num_samples=1000)\r\n```\r\nThanks :)\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport math\n\nfrom jax import jit\n\nimport numpyro\nfrom numpyro.compat.pyro import get_param_store\nfrom numpyro.infer import elbo, hmc, mcmc, svi\n\n\nclass HMC(hmc.HMC):\n    def __init__(\n        self,\n        model=None,\n        potential_fn=None,\n        step_size=1,\n        adapt_step_size=True,\n        adapt_mass_matrix=True,\n        full_mass=False,\n        use_multinomial_sampling=True,\n        transforms=None,\n        max_plate_nesting=None,\n        jit_compile=False,\n        jit_options=None,\n        ignore_jit_warnings=False,\n        trajectory_length=2 * math.pi,\n        target_accept_prob=0.8,\n    ):\n        super(HMC, self).__init__(\n            model=model,\n            potential_fn=potential_fn,\n            step_size=step_size,\n            adapt_step_size=adapt_step_size,\n            adapt_mass_matrix=adapt_mass_matrix,\n            dense_mass=full_mass,\n            target_accept_prob=target_accept_prob,\n            trajectory_length=trajectory_length,\n        )\n\n\nclass NUTS(hmc.NUTS):\n    def __init__(\n        self,\n        model=None,\n        potential_fn=None,\n        step_size=1,\n        adapt_step_size=True,\n        adapt_mass_matrix=True,\n        full_mass=False,\n        use_multinomial_sampling=True,\n        transforms=None,\n        max_plate_nesting=None,\n        jit_compile=False,\n        jit_options=None,\n        ignore_jit_warnings=False,\n        trajectory_length=2 * math.pi,\n        target_accept_prob=0.8,\n        max_tree_depth=10,\n    ):\n        if potential_fn is not None:\n            raise ValueError(\n                \"Only `model` argument is supported in generic module;\"\n                \" `potential_fn` is not supported.\"\n            )\n        super(NUTS, self).__init__(\n            model=model,\n            potential_fn=potential_fn,\n            step_size=step_size,\n            adapt_step_size=adapt_step_size,\n            adapt_mass_matrix=adapt_mass_matrix,\n            dense_mass=full_mass,\n            target_accept_prob=target_accept_prob,\n            trajectory_length=trajectory_length,\n            max_tree_depth=max_tree_depth,\n        )\n\n\nclass MCMC(object):\n    def __init__(\n        self,\n        kernel,\n        num_samples,\n        num_warmup=None,\n        initial_params=None,\n        num_chains=1,\n        hook_fn=None,\n        mp_context=None,\n        disable_progbar=False,\n        disable_validation=True,\n        transforms=None,\n    ):\n        if num_warmup is None:\n            num_warmup = num_samples\n        self._initial_params = initial_params\n        self._mcmc = mcmc.MCMC(\n            kernel,\n            num_warmup,\n            num_samples,\n            num_chains=num_chains,\n            progress_bar=(not disable_progbar),\n        )\n\n    def run(self, *args, rng_key=None, **kwargs):\n        if rng_key is None:\n            rng_key = numpyro.prng_key()\n        self._mcmc.run(rng_key, *args, init_params=self._initial_params, **kwargs)\n\n    def get_samples(self, num_samples=None, group_by_chain=False):\n        if num_samples is not None:\n            raise ValueError(\"`num_samples` arg unsupported in NumPyro.\")\n        return self._mcmc.get_samples(group_by_chain=group_by_chain)\n\n    def summary(self, prob=0.9):\n        self._mcmc.print_summary()\n\n\nclass SVI(svi.SVI):\n    def __init__(\n        self,\n        model,\n        guide,\n        optim,\n        loss,\n        loss_and_grads=None,\n        num_samples=10,\n        num_steps=0,\n        **kwargs\n    ):\n        super(SVI, self).__init__(model=model, guide=guide, optim=optim, loss=loss)\n        self.svi_state = None\n\n    def evaluate_loss(self, *args, **kwargs):\n        return self.evaluate(self.svi_state, *args, **kwargs)\n\n    def step(self, *args, rng_key=None, **kwargs):\n        if self.svi_state is None:\n            if rng_key is None:\n                rng_key = numpyro.prng_key()\n            self.svi_state = self.init(rng_key, *args, **kwargs)\n        try:\n            self.svi_state, loss = jit(self.update)(self.svi_state, *args, **kwargs)\n        except TypeError as e:\n            if \"not a valid JAX type\" in str(e):\n                raise TypeError(\n                    \"NumPyro backend requires args, kwargs to be arrays or tuples, \"\n                    \"dicts of arrays.\"\n                ) from e\n            else:\n                raise e\n        params = jit(super(SVI, self).get_params)(self.svi_state)\n        get_param_store().update(params)\n        return loss\n\n    def get_params(self):\n        return super(SVI, self).get_params(self.svi_state)\n\n\nclass Trace_ELBO(elbo.Trace_ELBO):\n    def __init__(\n        self,\n        num_particles=1,\n        max_plate_nesting=float(\"inf\"),\n        max_iarange_nesting=None,  # DEPRECATED\n        vectorize_particles=False,\n        strict_enumeration_warning=True,\n        ignore_jit_warnings=False,\n        jit_options=None,\n        retain_graph=None,\n        tail_adaptive_beta=-1.0,\n    ):\n        super(Trace_ELBO, self).__init__(num_particles=num_particles)\n\n\n# JIT is enabled by default\nJitTrace_ELBO = Trace_ELBO\n", "path": "numpyro/compat/infer.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport math\n\nfrom jax import jit\n\nimport numpyro\nfrom numpyro.compat.pyro import get_param_store\nfrom numpyro.infer import elbo, hmc, mcmc, svi\n\n\nclass HMC(hmc.HMC):\n    def __init__(\n        self,\n        model=None,\n        potential_fn=None,\n        step_size=1,\n        adapt_step_size=True,\n        adapt_mass_matrix=True,\n        full_mass=False,\n        use_multinomial_sampling=True,\n        transforms=None,\n        max_plate_nesting=None,\n        jit_compile=False,\n        jit_options=None,\n        ignore_jit_warnings=False,\n        trajectory_length=2 * math.pi,\n        target_accept_prob=0.8,\n    ):\n        super(HMC, self).__init__(\n            model=model,\n            potential_fn=potential_fn,\n            step_size=step_size,\n            adapt_step_size=adapt_step_size,\n            adapt_mass_matrix=adapt_mass_matrix,\n            dense_mass=full_mass,\n            target_accept_prob=target_accept_prob,\n            trajectory_length=trajectory_length,\n        )\n\n\nclass NUTS(hmc.NUTS):\n    def __init__(\n        self,\n        model=None,\n        potential_fn=None,\n        step_size=1,\n        adapt_step_size=True,\n        adapt_mass_matrix=True,\n        full_mass=False,\n        use_multinomial_sampling=True,\n        transforms=None,\n        max_plate_nesting=None,\n        jit_compile=False,\n        jit_options=None,\n        ignore_jit_warnings=False,\n        trajectory_length=2 * math.pi,\n        target_accept_prob=0.8,\n        max_tree_depth=10,\n    ):\n        if potential_fn is not None:\n            raise ValueError(\n                \"Only `model` argument is supported in generic module;\"\n                \" `potential_fn` is not supported.\"\n            )\n        super(NUTS, self).__init__(\n            model=model,\n            potential_fn=potential_fn,\n            step_size=step_size,\n            adapt_step_size=adapt_step_size,\n            adapt_mass_matrix=adapt_mass_matrix,\n            dense_mass=full_mass,\n            target_accept_prob=target_accept_prob,\n            trajectory_length=trajectory_length,\n            max_tree_depth=max_tree_depth,\n        )\n\n\nclass MCMC(object):\n    def __init__(\n        self,\n        kernel,\n        num_samples,\n        warmup_steps=None,\n        initial_params=None,\n        num_chains=1,\n        hook_fn=None,\n        mp_context=None,\n        disable_progbar=False,\n        disable_validation=True,\n        transforms=None,\n    ):\n        if warmup_steps is None:\n            warmup_steps = num_samples\n        self._initial_params = initial_params\n        self._mcmc = mcmc.MCMC(\n            kernel,\n            num_warmup=warmup_steps,\n            num_samples=num_samples,\n            num_chains=num_chains,\n            progress_bar=(not disable_progbar),\n        )\n\n    def run(self, *args, rng_key=None, **kwargs):\n        if rng_key is None:\n            rng_key = numpyro.prng_key()\n        self._mcmc.run(rng_key, *args, init_params=self._initial_params, **kwargs)\n\n    def get_samples(self, num_samples=None, group_by_chain=False):\n        if num_samples is not None:\n            raise ValueError(\"`num_samples` arg unsupported in NumPyro.\")\n        return self._mcmc.get_samples(group_by_chain=group_by_chain)\n\n    def summary(self, prob=0.9):\n        self._mcmc.print_summary()\n\n\nclass SVI(svi.SVI):\n    def __init__(\n        self,\n        model,\n        guide,\n        optim,\n        loss,\n        loss_and_grads=None,\n        num_samples=10,\n        num_steps=0,\n        **kwargs\n    ):\n        super(SVI, self).__init__(model=model, guide=guide, optim=optim, loss=loss)\n        self.svi_state = None\n\n    def evaluate_loss(self, *args, **kwargs):\n        return self.evaluate(self.svi_state, *args, **kwargs)\n\n    def step(self, *args, rng_key=None, **kwargs):\n        if self.svi_state is None:\n            if rng_key is None:\n                rng_key = numpyro.prng_key()\n            self.svi_state = self.init(rng_key, *args, **kwargs)\n        try:\n            self.svi_state, loss = jit(self.update)(self.svi_state, *args, **kwargs)\n        except TypeError as e:\n            if \"not a valid JAX type\" in str(e):\n                raise TypeError(\n                    \"NumPyro backend requires args, kwargs to be arrays or tuples, \"\n                    \"dicts of arrays.\"\n                ) from e\n            else:\n                raise e\n        params = jit(super(SVI, self).get_params)(self.svi_state)\n        get_param_store().update(params)\n        return loss\n\n    def get_params(self):\n        return super(SVI, self).get_params(self.svi_state)\n\n\nclass Trace_ELBO(elbo.Trace_ELBO):\n    def __init__(\n        self,\n        num_particles=1,\n        max_plate_nesting=float(\"inf\"),\n        max_iarange_nesting=None,  # DEPRECATED\n        vectorize_particles=False,\n        strict_enumeration_warning=True,\n        ignore_jit_warnings=False,\n        jit_options=None,\n        retain_graph=None,\n        tail_adaptive_beta=-1.0,\n    ):\n        super(Trace_ELBO, self).__init__(num_particles=num_particles)\n\n\n# JIT is enabled by default\nJitTrace_ELBO = Trace_ELBO\n", "path": "numpyro/compat/infer.py"}]}