prompt: 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()
  
  ```
  
golden_diff: 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),
           )
  
verification_info: {"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"}]}