internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: vllm-project/vllm

Issue:
[Bug]: Tokenizer setter of LLM without CachedTokenizer adapter
### Your current environment

same to the issue [#5206](https://github.com/vllm-project/vllm/issues/5206)


### 🐛 Describe the bug

As the basic reason of the issue reported by [#5206](https://github.com/vllm-project/vllm/issues/5206), the tokenizer setter of the LLM will override the cached tokenizer inited by llm_engine. 
```python
    def set_tokenizer(
        self,
        tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
    ) -> None:
        self.llm_engine.tokenizer.tokenizer = tokenizer
```
Thus, each time the ```len(tokenizer)``` is called, row ```__len__``` is called rather than the CachedTokenizer's.
To fixed this problem of this issue, the ```get_cached_tokenizer``` adapter should be applied by the tokenizer setter of LLM.


Discussion:

Resolving patch:
diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py
--- a/vllm/entrypoints/llm.py
+++ b/vllm/entrypoints/llm.py
@@ -14,6 +14,7 @@
 from vllm.outputs import EmbeddingRequestOutput, RequestOutput
 from vllm.pooling_params import PoolingParams
 from vllm.sampling_params import SamplingParams
+from vllm.transformers_utils.tokenizer import get_cached_tokenizer
 from vllm.usage.usage_lib import UsageContext
 from vllm.utils import Counter, deprecate_kwargs
 
@@ -152,7 +153,14 @@ def set_tokenizer(
         self,
         tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
     ) -> None:
-        self.llm_engine.tokenizer.tokenizer = tokenizer
+        # While CachedTokenizer is dynamic, have no choice but
+        # compare class name. Misjudgment will arise from
+        # user-defined tokenizer started with 'Cached'
+        if tokenizer.__class__.__name__.startswith("Cached"):
+            self.llm_engine.tokenizer.tokenizer = tokenizer
+        else:
+            self.llm_engine.tokenizer.tokenizer = get_cached_tokenizer(
+                tokenizer)
 
     @overload  # LEGACY: single (prompt + optional token ids)
     def generate(


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Make HMCState picklable
Currently, we use `laxtuple` to make HMCState. This way however make it not possible to pickle HMCState (unless we transform it to a dict/tuple/namedtuple).


Discussion:

Resolving patch:
diff --git a/numpyro/mcmc.py b/numpyro/mcmc.py
--- a/numpyro/mcmc.py
+++ b/numpyro/mcmc.py
@@ -1,5 +1,6 @@
 import math
 import os
+from collections import namedtuple
 
 import tqdm
 
@@ -7,13 +8,24 @@
 from jax import jit, partial, random
 from jax.flatten_util import ravel_pytree
 from jax.random import PRNGKey
+from jax.tree_util import register_pytree_node
 
 import numpyro.distributions as dist
 from numpyro.hmc_util import IntegratorState, build_tree, find_reasonable_step_size, velocity_verlet, warmup_adapter
-from numpyro.util import cond, fori_loop, laxtuple
+from numpyro.util import cond, fori_loop
 
-HMCState = laxtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',
-                                 'step_size', 'inverse_mass_matrix', 'rng'])
+HMCState = namedtuple('HMCState', ['z', 'z_grad', 'potential_energy', 'num_steps', 'accept_prob',
+                                   'step_size', 'inverse_mass_matrix', 'rng'])
+
+
+register_pytree_node(
+    HMCState,
+    lambda xs: (tuple(xs), None),
+    lambda _, xs: HMCState(*xs)
+)
+
+
+HMCState.update = HMCState._replace
 
 
 def _get_num_steps(step_size, trajectory_length):


Test patch:
Repository: pyro-ppl/numpyro

Issue:
HMC adaptation stuck at warmup phase due to step_size -> nan
A script to replicate this issue
```
import jax.numpy as np
from jax.scipy.special import logsumexp
from numpyro.mcmc import mcmc

def dual_moon_pe(x):
    term1 = 0.5 * ((np.linalg.norm(x, axis=-1) - 2) / 0.4) ** 2
    term2 = -0.5 * ((x[..., :1] + np.array([-2., 2.])) / 0.6) ** 2
    return term1 - logsumexp(term2, axis=-1)

mcmc(1000, 1000, init_params=np.array([2., 0.]), potential_fn=dual_moon_pe)
```


Discussion:
This seems to run fine on my system, even for much higher warmup steps, but gives poor results, as you might expect.

Update - I can replicate it with a different rng. 

Resolving patch:
diff --git a/numpyro/hmc_util.py b/numpyro/hmc_util.py
--- a/numpyro/hmc_util.py
+++ b/numpyro/hmc_util.py
@@ -248,6 +248,7 @@ def _body_fn(state):
         # case for a diverging trajectory (e.g. in the case of evaluating log prob
         # of a value simulated using a large step size for a constrained sample site).
         step_size = (2.0 ** direction) * step_size
+        step_size = np.clip(step_size, a_min=np.finfo(step_size.dtype).tiny)
         r = momentum_generator(inverse_mass_matrix, rng_momentum)
         _, r_new, potential_energy_new, _ = vv_update(step_size,
                                                       inverse_mass_matrix,


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Clean up docs/source/conf.py file
I think we can change the names `Numpyro` -> `NumPyro` there, but I am not sure if the changes will affect the build of the website. So I make this issue.

cc @jpchen @neerajprad 


Discussion:
Please go ahead! I think this was just a typo to begin with.

Resolving patch:
diff --git a/docs/source/conf.py b/docs/source/conf.py
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -33,7 +33,7 @@
 
 # -- Project information -----------------------------------------------------
 
-project = u'Numpyro'
+project = u'NumPyro'
 copyright = u'2019, Uber Technologies, Inc'
 author = u'Uber AI Labs'
 
@@ -165,7 +165,7 @@
 # (source start file, target name, title,
 #  author, documentclass [howto, manual, or own class]).
 latex_documents = [
-    (master_doc, 'Numpyro.tex', u'Numpyro Documentation', u'Uber AI Labs', 'manual'),
+    (master_doc, 'NumPyro.tex', u'NumPyro Documentation', u'Uber AI Labs', 'manual'),
 ]
 
 # -- Options for manual page output ------------------------------------------
@@ -173,7 +173,7 @@
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
 man_pages = [
-    (master_doc, 'Numpyro', u'Numpyro Documentation',
+    (master_doc, 'NumPyro', u'NumPyro Documentation',
      [author], 1)
 ]
 
@@ -183,8 +183,8 @@
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
-    (master_doc, 'Numpyro', u'Numpyro Documentation',
-     author, 'Numpyro', 'Pyro PPL on Numpy',
+    (master_doc, 'NumPyro', u'NumPyro Documentation',
+     author, 'NumPyro', 'Pyro PPL on Numpy',
      'Miscellaneous'),
 ]
 


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Regression in fori_collect when progress bar disabled
I am finding that with progress bar disabled, `fori_collect` is not returning the same results for the sparse regression example:
 - [x] It fetches the wrong `init_state` and runs adaptation twice.
 - [ ] Results returned when progress bar is enabled/disabled are different.
 - [x] When progress bar is disabled, the time taken is almost twice. This seems to be a regression from earlier when disabling the progress bar was much faster. 


Discussion:

Resolving patch:
diff --git a/numpyro/util.py b/numpyro/util.py
--- a/numpyro/util.py
+++ b/numpyro/util.py
@@ -181,11 +181,11 @@ def fori_collect(lower, upper, body_fun, init_val, transform=identity,
 
         def _body_fn(i, vals):
             val, collection, start_state = vals
-            val = body_fun(val)
-            i = np.where(i >= lower, i - lower, 0)
-            start_state = lax.cond(i == lower-1,
-                                   start_state, lambda _: val,
+            val = jit(body_fun)(val)
+            start_state = lax.cond(i < lower,
+                                   val, lambda x: x,
                                    start_state, lambda x: x)
+            i = np.where(i >= lower, i - lower, 0)
             collection = ops.index_update(collection, i, ravel_fn(val))
             return val, collection, start_state
 


Test patch:
diff --git a/test/test_util.py b/test/test_util.py
--- a/test/test_util.py
+++ b/test/test_util.py
@@ -20,6 +20,22 @@ def f(x):
     check_eq(actual_tree, expected_tree)
 
 
+@pytest.mark.parametrize('progbar', [False, True])
+def test_fori_collect_return_init(progbar):
+    def f(x):
+        x['i'] = x['i'] + 1
+        return x
+
+    tree, init_state = fori_collect(2, 4, f, {'i': 0},
+                                    transform=lambda a: {'i': a['i']},
+                                    return_init_state=True,
+                                    progbar=progbar)
+    expected_tree = {'i': np.array([3, 4])}
+    expected_init_state = {'i': np.array(2)}
+    check_eq(init_state, expected_init_state)
+    check_eq(tree, expected_tree)
+
+
 @pytest.mark.parametrize('pytree', [
     {'a': np.array(0.), 'b': np.array([[1., 2.], [3., 4.]])},
     {'a': np.array(0), 'b': np.array([[1, 2], [3, 4]])},
Repository: pyro-ppl/numpyro

Issue:
Model compilation happens twice with the latest JAX version
With the latest jax version, many models are compiled twice - once at the beginning of warmup and then again when sampling starts. The reason for this is the introduction of weak types in JAX (all python scalars have a weak type which might be promoted later depending on the dtype of other arrays). 

As such, the arguments to `sample_fn` at the beginning of both the stages may differ in this aspect, i.e. arguments passed by the user may have weak types but those that are used to begin sampling don't, which results in a cache miss. The full impact of this still needs to be assessed. It will probably not matter for small models.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -6,7 +6,7 @@
 import os
 import warnings
 
-from jax import jit, lax, partial, pmap, random, vmap, device_get
+from jax import jit, lax, partial, pmap, random, vmap, device_get, device_put
 from jax.core import Tracer
 from jax.dtypes import canonicalize_dtype
 from jax.flatten_util import ravel_pytree
@@ -254,7 +254,7 @@ def init_kernel(init_params,
         energy = kinetic_fn(wa_state.inverse_mass_matrix, vv_state.r)
         hmc_state = HMCState(0, vv_state.z, vv_state.z_grad, vv_state.potential_energy, energy,
                              0, 0., 0., False, wa_state, rng_key_hmc)
-        return hmc_state
+        return device_put(hmc_state)
 
     def _hmc_next(step_size, inverse_mass_matrix, vv_state,
                   model_args, model_kwargs, rng_key):


Test patch:
Repository: pyro-ppl/numpyro

Issue:
mcmc progbar does not display correctly running phase
The last change does not update the logic of progressbar. So it displays `warmup`, `sample` incorrectly when we run `.warmup` separately from `.run`.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -754,7 +754,6 @@ def _get_cached_init_state(self, rng_key, args, kwargs):
             return None
 
     def _single_chain_mcmc(self, rng_key, init_state, init_params, args, kwargs, collect_fields=('z',)):
-        num_warmup = self.num_warmup
         if init_state is None:
             init_state = self.sampler.init(rng_key, self.num_warmup, init_params,
                                            model_args=args, model_kwargs=kwargs)
@@ -762,15 +761,19 @@ def _single_chain_mcmc(self, rng_key, init_state, init_params, args, kwargs, col
             self.constrain_fn = self.sampler.constrain_fn(args, kwargs)
         diagnostics = lambda x: get_diagnostics_str(x[0]) if rng_key.ndim == 1 else None   # noqa: E731
         init_val = (init_state, args, kwargs) if self._jit_model_args else (init_state,)
-        collect_vals = fori_collect(self._collection_params["lower"],
-                                    self._collection_params["upper"],
+        lower_idx = self._collection_params["lower"]
+        upper_idx = self._collection_params["upper"]
+
+        collect_vals = fori_collect(lower_idx,
+                                    upper_idx,
                                     self._get_cached_fn(),
                                     init_val,
                                     transform=_collect_fn(collect_fields),
                                     progbar=self.progress_bar,
                                     return_last_val=True,
                                     collection_size=self._collection_params["collection_size"],
-                                    progbar_desc=functools.partial(get_progbar_desc_str, num_warmup),
+                                    progbar_desc=functools.partial(get_progbar_desc_str,
+                                                                   num_warmup=lower_idx),
                                     diagnostics_fn=diagnostics)
         states, last_val = collect_vals
         # Get first argument of type `HMCState`


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Use ExpandedDistribution in numpyro.plate 
Since #574, I think we should be able to switch our existing logic to use `ExpandedDistribution` instead of our current logic. This will align the underlying implementation with that of Pyro and possibly aid with funsor integration.


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -311,6 +311,7 @@ class ExpandedDistribution(Distribution):
 
     def __init__(self, base_dist, batch_shape=()):
         if isinstance(base_dist, ExpandedDistribution):
+            batch_shape = self._broadcast_shape(base_dist.batch_shape, batch_shape)
             base_dist = base_dist.base_dist
         self.base_dist = base_dist
         super().__init__(base_dist.batch_shape, base_dist.event_shape)
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -263,19 +263,16 @@ def process_message(self, msg):
         frame = CondIndepStackFrame(self.name, self.dim, self.subsample_size)
         cond_indep_stack.append(frame)
         expected_shape = self._get_batch_shape(cond_indep_stack)
-        dist_batch_shape = msg['fn'].batch_shape if msg['type'] == 'sample' else ()
-        overlap_idx = max(len(expected_shape) - len(dist_batch_shape), 0)
-        trailing_shape = expected_shape[overlap_idx:]
-        # e.g. distribution with batch shape (1, 5) cannot be broadcast to (5, 5)
-        broadcast_shape = lax.broadcast_shapes(trailing_shape, dist_batch_shape)
-        if broadcast_shape != dist_batch_shape:
-            raise ValueError('Distribution batch shape = {} cannot be broadcast up to {}. '
-                             'Consider using unbatched distributions.'
-                             .format(dist_batch_shape, broadcast_shape))
-        batch_shape = expected_shape[:overlap_idx]
-        if 'sample_shape' in msg['kwargs']:
-            batch_shape = lax.broadcast_shapes(msg['kwargs']['sample_shape'], batch_shape)
-        msg['kwargs']['sample_shape'] = batch_shape
+        if msg['type'] == 'sample':
+            dist_batch_shape = msg['fn'].batch_shape
+            if 'sample_shape' in msg['kwargs']:
+                dist_batch_shape = msg['kwargs']['sample_shape'] + dist_batch_shape
+                msg['kwargs']['sample_shape'] = ()
+            overlap_idx = max(len(expected_shape) - len(dist_batch_shape), 0)
+            trailing_shape = expected_shape[overlap_idx:]
+            broadcast_shape = lax.broadcast_shapes(trailing_shape, dist_batch_shape)
+            batch_shape = expected_shape[:overlap_idx] + broadcast_shape
+            msg['fn'] = msg['fn'].expand(batch_shape)
         if self.size != self.subsample_size:
             scale = 1. if msg['scale'] is None else msg['scale']
             msg['scale'] = scale * self.size / self.subsample_size


Test patch:
Repository: pyro-ppl/numpyro

Issue:
Support `hide` argument in `block`
If we want to hide some sites, using `hide_fn=lambda site: site['name'] == 'x'` is a bit inconvenient. We can mimic Pyro by allowing an extra `hide` argument, so users can just simply use `hide=['x']`.


Discussion:

Resolving patch:
diff --git a/numpyro/contrib/distributions/distribution.py b/numpyro/contrib/distributions/distribution.py
--- a/numpyro/contrib/distributions/distribution.py
+++ b/numpyro/contrib/distributions/distribution.py
@@ -9,10 +9,25 @@
 # Copyright (c) 2003-2019 SciPy Developers.
 # All rights reserved.
 from contextlib import contextmanager
+try:
+    from scipy._lib._util import getargspec_no_self
+except ImportError:
+    from collections import namedtuple
+    from scipy._lib._util import getfullargspec_no_self
 
-from scipy._lib._util import getargspec_no_self
+    ArgSpec = namedtuple('ArgSpec', ['args', 'varargs', 'keywords', 'defaults'])
+
+    def getargspec_no_self(obj):
+        return ArgSpec(*getfullargspec_no_self(obj)[:4])
 import scipy.stats as osp_stats
-from scipy.stats._distn_infrastructure import instancemethod, rv_frozen, rv_generic
+from scipy.stats._distn_infrastructure import rv_frozen, rv_generic
+try:
+    from scipy.stats._distn_infrastructure import instancemethod
+except ImportError:
+    import types
+
+    def instancemethod(func, obj, cls):
+        return types.MethodType(func, obj)
 
 from jax import lax
 import jax.numpy as jnp
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -220,8 +220,13 @@ class block(Messenger):
        >>> assert 'a' not in trace_block_a
        >>> assert 'b' in trace_block_a
     """
-    def __init__(self, fn=None, hide_fn=lambda msg: True):
-        self.hide_fn = hide_fn
+    def __init__(self, fn=None, hide_fn=None, hide=None):
+        if hide_fn is not None:
+            self.hide_fn = hide_fn
+        elif hide is not None:
+            self.hide_fn = lambda msg: msg.get('name') in hide
+        else:
+            self.hide_fn = lambda msg: True
         super(block, self).__init__(fn)
 
     def process_message(self, msg):


Test patch:
diff --git a/test/test_handlers.py b/test/test_handlers.py
--- a/test/test_handlers.py
+++ b/test/test_handlers.py
@@ -242,6 +242,14 @@ def guide():
     assert x.shape == shape
 
 
+def test_block():
+    with handlers.trace() as trace:
+        with handlers.block(hide=['x']):
+            with handlers.seed(rng_seed=0):
+                numpyro.sample('x', dist.Normal())
+    assert 'x' not in trace
+
+
 def test_scope():
     def fn():
         return numpyro.sample('x', dist.Normal())
Repository: pyro-ppl/numpyro

Issue:
Cannot reshape the output of chained MCMC
This issue is created to make sure that https://github.com/google/jax/issues/3597 is resolved before we release the next version.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -7,7 +7,7 @@
 import os
 import warnings
 
-from jax import jit, lax, pmap, random
+from jax import device_put, jit, lax, pmap, random
 from jax.core import Tracer
 from jax.interpreters.xla import DeviceArray
 from jax.lib import xla_bridge
@@ -453,6 +453,8 @@ def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
                     states, last_state = lax.map(partial_map_fn, map_args)
             elif chain_method == 'parallel':
                 states, last_state = pmap(partial_map_fn)(map_args)
+                # TODO: remove when https://github.com/google/jax/issues/3597 is resolved
+                states = device_put(states)
             else:
                 assert chain_method == 'vectorized'
                 states, last_state = partial_map_fn(map_args)


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -7,7 +7,7 @@
 from numpy.testing import assert_allclose
 import pytest
 
-from jax import jit, pmap, random, vmap
+from jax import device_get, jit, pmap, random, vmap
 from jax.lib import xla_bridge
 import jax.numpy as jnp
 from jax.scipy.special import logit
@@ -380,6 +380,9 @@ def model(labels):
     assert samples['coefs'].shape[:2] == (num_chains, num_samples)
     assert_allclose(jnp.mean(samples_flat['coefs'], 0), true_coefs, atol=0.21)
 
+    # test if reshape works
+    device_get(samples_flat['coefs'].reshape(-1))
+
 
 @pytest.mark.parametrize('kernel_cls', [HMC, NUTS])
 @pytest.mark.parametrize('chain_method', [
Repository: pyro-ppl/numpyro

Issue:
Add a better error message for how to debug "Cannot find valid initial parameters" issue
Recently, several users got this issue but it is not easy to diagnose the issue just by looking at the `model`. We should add a better error message, e.g. to mention about `numpyro.validation_enabled()` utility, which tells us at which latent site, we get invalid parameters/values.


Discussion:
We can also run the model with `validation_enabled` once if we can't find a valid initial parameters.

Resolving patch:
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -141,7 +141,8 @@ def __init__(self, batch_shape=(), event_shape=(), validate_args=None):
                 is_valid = jnp.all(constraint(getattr(self, param)))
                 if not_jax_tracer(is_valid):
                     if not is_valid:
-                        raise ValueError("The parameter {} has invalid values".format(param))
+                        raise ValueError("{} distribution got invalid {} parameter.".format(
+                            self.__class__.__name__, param))
         super(Distribution, self).__init__()
 
     @property
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -427,6 +427,21 @@ def initialize_model(rng_key, model,
 
     if not_jax_tracer(is_valid):
         if device_get(~jnp.all(is_valid)):
+            with numpyro.validation_enabled(), trace() as tr:
+                # validate parameters
+                substituted_model(*model_args, **model_kwargs)
+                # validate values
+                for site in tr.values():
+                    if site['type'] == 'sample':
+                        with warnings.catch_warnings(record=True) as ws:
+                            site['fn']._validate_sample(site['value'])
+                        if len(ws) > 0:
+                            for w in ws:
+                                # at site information to the warning message
+                                w.message.args = ("Site {}: {}".format(site["name"], w.message.args[0]),) \
+                                    + w.message.args[1:]
+                                warnings.showwarning(w.message, w.category, w.filename, w.lineno,
+                                                     file=w.file, line=w.line)
             raise RuntimeError("Cannot find valid initial parameters. Please check your model again.")
     return ModelInfo(ParamInfo(init_params, pe, grad), potential_fn, postprocess_fn, model_trace)
 


Test patch: