Repository: pyro-ppl/numpyro

Issue:
`init_to_uniform` sometimes leads to mangled chains
I'm not quite sure if this a bug or expected behavior, but the behavior was at least unexpected and undesirable to me.

Note that even with the `init_to_uniform` `init_strategy`, not every invocation of `run` results in exactly one mangled chain—sometimes there are zero, sometimes there are multiple.

Would it make sense to think about the default initialization strategy used by the library? `pymc3`'s defaults seem to handle this correctly for an equivalent model.

```python
import numpy as np

import numpyro as ny
from numpyro.infer import MCMC, NUTS
import numpyro.distributions as dist
import jax.numpy as jnp
from jax import random
import jax

import pandas as pd
import arviz as az

az.style.use("arviz-darkgrid")

az.rcParams["plot.max_subplots"] = 1000

# pd.options.display.max_rows = 1000
pd.set_option('display.max_colwidth', None)
pd.set_option('display.max_columns', None)
pd.set_option('display.width', 2000)

ny.util.set_host_device_count(6)
```


```python
def model(data):
    width = data.shape[1]
    length = data.shape[0]

    corr = ny.sample("corr", dist.LKJ(width, concentration=1))
    std_devs = ny.sample("std_devs", dist.HalfCauchy(jnp.ones(width)))
    chol = ny.deterministic("chol", jnp.linalg.cholesky(jnp.matmul(jnp.matmul(jnp.diag(std_devs), corr), jnp.diag(std_devs))))
    chol_l = ny.deterministic("chol_l", chol[jnp.tril_indices(width)])
    
    mu_reparam = ny.sample("mu_reparam", dist.MultivariateNormal(loc=jnp.zeros(width), covariance_matrix=jnp.identity(width)))
    mu = ny.deterministic("mu", jnp.matmul(chol,  mu_reparam))
    
    with ny.plate("data", length):
        returns = ny.sample("returns", dist.MultivariateNormal(loc=mu, scale_tril=chol), obs=data)
```


```python
def run(chain_method, data, adapt_step_size, step_size, init_strategy):
    mcmc = MCMC(NUTS(model, adapt_step_size=adapt_step_size, step_size=step_size, init_strategy=init_strategy), num_warmup=1000, num_samples=1000, num_chains=4, progress_bar=True, chain_method=chain_method)
    mcmc.run(random.PRNGKey(0), data=data, extra_fields=["diverging", "potential_energy"])

    idata = az.from_numpyro(mcmc)
    print(az.summary(mcmc, var_names=["mu", "chol_l"]))
    az.plot_rank(idata, var_names=["chol_l", "mu"])
    # az.plot_autocorr(idata, var_names=["chol_l", "mu"]);
```


```python
run(chain_method="sequential", data=np.random.normal(loc=1, scale=5, size=(500,7)), adapt_step_size=True, step_size=1, init_strategy=ny.infer.initialization.init_to_median)
run(chain_method="sequential", data=np.random.normal(loc=1, scale=5, size=(500,7)), adapt_step_size=True, step_size=1, init_strategy=ny.infer.initialization.init_to_uniform)
```

    sample: 100%|██████████| 2000/2000 [00:09<00:00, 208.90it/s, 7 steps of size 5.28e-01. acc. prob=0.86] 
    sample: 100%|██████████| 2000/2000 [00:01<00:00, 1196.41it/s, 7 steps of size 5.42e-01. acc. prob=0.85]
    sample: 100%|██████████| 2000/2000 [00:02<00:00, 962.71it/s, 7 steps of size 4.68e-01. acc. prob=0.90] 
    sample: 100%|██████████| 2000/2000 [00:01<00:00, 1074.71it/s, 7 steps of size 5.05e-01. acc. prob=0.88]


                 mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
    mu[0]       1.053  0.227   0.623    1.475      0.003    0.002    6797.0    2961.0    1.0
    mu[1]       1.221  0.236   0.760    1.658      0.003    0.002    6368.0    3236.0    1.0
    mu[2]       0.561  0.220   0.133    0.956      0.003    0.002    6863.0    3168.0    1.0
    mu[3]       1.088  0.226   0.663    1.510      0.003    0.002    6980.0    3307.0    1.0
    mu[4]       1.207  0.222   0.816    1.643      0.003    0.002    7789.0    2926.0    1.0
    mu[5]       0.700  0.219   0.295    1.111      0.003    0.002    6463.0    2785.0    1.0
    mu[6]       1.248  0.208   0.856    1.632      0.002    0.002    7832.0    3265.0    1.0
    chol_l[0]   5.078  0.167   4.756    5.380      0.002    0.001    7115.0    2933.0    1.0
    chol_l[1]   0.152  0.240  -0.301    0.578      0.003    0.004    6097.0    2433.0    1.0
    chol_l[2]   5.181  0.170   4.880    5.512      0.002    0.002    6060.0    2878.0    1.0
    chol_l[3]  -0.541  0.222  -0.971   -0.144      0.003    0.002    6102.0    3177.0    1.0
    chol_l[4]  -0.067  0.214  -0.465    0.349      0.003    0.003    5438.0    2997.0    1.0
    chol_l[5]   4.819  0.152   4.543    5.105      0.002    0.001    5801.0    3452.0    1.0
    chol_l[6]  -0.284  0.227  -0.705    0.154      0.003    0.003    5767.0    3285.0    1.0
    chol_l[7]  -0.159  0.230  -0.564    0.297      0.003    0.003    5632.0    3035.0    1.0
    chol_l[8]  -0.316  0.223  -0.732    0.105      0.003    0.002    5990.0    3035.0    1.0
    chol_l[9]   5.061  0.159   4.754    5.352      0.002    0.001    7467.0    2756.0    1.0
    chol_l[10] -0.447  0.217  -0.826   -0.014      0.003    0.002    5970.0    3070.0    1.0
    chol_l[11]  0.231  0.217  -0.190    0.625      0.003    0.002    6696.0    3573.0    1.0
    chol_l[12] -0.234  0.212  -0.628    0.159      0.003    0.003    6447.0    2938.0    1.0
    chol_l[13]  0.393  0.212   0.008    0.815      0.003    0.002    6304.0    3365.0    1.0
    chol_l[14]  4.808  0.151   4.530    5.092      0.002    0.001    9119.0    3033.0    1.0
    chol_l[15] -0.236  0.214  -0.643    0.137      0.003    0.003    5638.0    3106.0    1.0
    chol_l[16] -0.050  0.215  -0.457    0.339      0.003    0.003    5248.0    3141.0    1.0
    chol_l[17] -0.166  0.221  -0.568    0.250      0.003    0.003    6806.0    3327.0    1.0
    chol_l[18] -0.180  0.215  -0.583    0.216      0.003    0.003    6921.0    3229.0    1.0
    chol_l[19]  0.522  0.216   0.127    0.922      0.003    0.002    6544.0    2983.0    1.0
    chol_l[20]  4.841  0.151   4.562    5.125      0.002    0.001    6960.0    3216.0    1.0
    chol_l[21]  0.179  0.212  -0.215    0.570      0.002    0.003    7377.0    3212.0    1.0
    chol_l[22]  0.232  0.209  -0.157    0.626      0.003    0.003    6552.0    3162.0    1.0
    chol_l[23] -0.200  0.203  -0.612    0.163      0.002    0.002    7229.0    3132.0    1.0
    chol_l[24] -0.025  0.201  -0.400    0.346      0.002    0.003    7460.0    3388.0    1.0
    chol_l[25] -0.103  0.215  -0.499    0.311      0.003    0.003    6562.0    2818.0    1.0
    chol_l[26] -0.150  0.202  -0.508    0.239      0.003    0.003    6510.0    2922.0    1.0
    chol_l[27]  4.594  0.145   4.334    4.878      0.002    0.001    6160.0    2827.0    1.0


    sample: 100%|██████████| 2000/2000 [00:09<00:00, 214.14it/s, 7 steps of size 4.84e-01. acc. prob=0.88] 
    sample: 100%|██████████| 2000/2000 [02:08<00:00, 15.52it/s, 1023 steps of size 3.37e-09. acc. prob=0.85]
    sample: 100%|██████████| 2000/2000 [00:01<00:00, 1101.96it/s, 7 steps of size 5.14e-01. acc. prob=0.87]
    sample: 100%|██████████| 2000/2000 [00:02<00:00, 829.46it/s, 7 steps of size 4.97e-01. acc. prob=0.88] 


                 mean     sd  hdi_3%  hdi_97%  mcse_mean  mcse_sd  ess_bulk  ess_tail  r_hat
    mu[0]       0.590  0.701  -0.580    1.279      0.338    0.257       7.0       4.0   1.52
    mu[1]       0.730  0.415   0.100    1.270      0.183    0.138       7.0       4.0   1.52
    mu[2]       1.862  1.168   0.871    3.855      0.574    0.444       7.0    3263.0   1.52
    mu[3]       0.722  0.213   0.314    1.150      0.005    0.011    4567.0    1812.0   1.53
    mu[4]      -0.980  2.759  -5.745    0.930      1.371    1.052       7.0       4.0   1.53
    mu[5]       0.237  1.317  -2.017    1.312      0.649    0.496       7.0       4.0   1.53
    mu[6]       1.253  0.769   0.499    2.539      0.371    0.290       7.0    1565.0   1.53
    chol_l[0]   3.776  1.982   0.352    5.146      0.985    0.754       7.0       4.0   1.52
    chol_l[1]  -0.367  0.249  -0.835   -0.031      0.073    0.053      14.0    2105.0   1.21
    chol_l[2]   3.779  2.129   0.099    5.243      1.058    0.810       7.0       4.0   1.52
    chol_l[3]  -1.402  1.601  -4.152   -0.152      0.791    0.610       7.0       4.0   1.53
    chol_l[4]   0.795  0.800   0.030    2.136      0.387    0.304       7.0    2703.0   1.52
    chol_l[5]   4.111  1.717   1.147    5.335      0.853    0.652       7.0       4.0   1.52
    chol_l[6]   0.017  0.282  -0.329    0.547      0.094    0.069      11.0       5.0   1.28
    chol_l[7]  -0.477  0.272  -1.022   -0.170      0.089    0.065      11.0    2016.0   1.26
    chol_l[8]   0.305  0.300  -0.069    0.805      0.110    0.081       9.0       4.0   1.36
    chol_l[9]   4.100  2.332   0.069    5.688      1.159    0.887       7.0       4.0   1.52
    chol_l[10]  0.585  1.257  -0.449    2.735      0.619    0.479       7.0    2282.0   1.52
    chol_l[11]  0.167  0.325  -0.280    0.650      0.131    0.097       8.0       4.0   1.47
    chol_l[12]  0.211  0.196  -0.231    0.549      0.011    0.008     721.0    2244.0   1.52
    chol_l[13] -0.362  0.205  -0.722    0.070      0.031    0.022      54.0    2456.0   1.10
    chol_l[14]  4.001  1.824   0.852    5.275      0.906    0.693       7.0       4.0   1.53
    chol_l[15]  0.586  0.657  -0.102    1.670      0.313    0.249       7.0    2691.0   1.53
    chol_l[16]  0.609  1.175  -0.388    2.614      0.577    0.448       7.0    2441.0   1.52
    chol_l[17]  0.283  0.280  -0.198    0.629      0.098    0.083      10.0    2714.0   1.33
    chol_l[18] -0.075  0.201  -0.462    0.342      0.003    0.017    6376.0    2331.0   1.53
    chol_l[19] -0.336  1.165  -2.324    0.652      0.572    0.443       7.0       4.0   1.53
    chol_l[20]  4.119  1.876   0.879    5.426      0.932    0.713       7.0       4.0   1.53
    chol_l[21] -0.258  0.790  -1.581    0.506      0.382    0.300       7.0       4.0   1.52
    chol_l[22]  0.200  0.238  -0.262    0.522      0.071    0.053      13.0    1944.0   1.22
    chol_l[23] -0.246  0.210  -0.688    0.107      0.040    0.029      35.0    2384.0   1.08
    chol_l[24]  0.118  0.216  -0.319    0.447      0.051    0.037      23.0    2390.0   1.11
    chol_l[25] -0.223  0.213  -0.678    0.137      0.027    0.027      79.0    1940.0   1.14
    chol_l[26]  0.062  0.238  -0.244    0.503      0.074    0.054      11.0       5.0   1.26
    chol_l[27]  3.760  2.131   0.078    5.216      1.059    0.811       7.0       4.0   1.53



    
![output_3_4](https://user-images.githubusercontent.com/9491942/130694716-9e065b56-0d5e-4177-9c5c-8ed9eadab869.png)
![output_3_5](https://user-images.githubusercontent.com/9491942/130694753-8457078d-8e74-4860-a61f-d4e47906075b.png)




Discussion:
Interesting! I think `init_to_uniform` is bad if it provides initial values in a low-density area of the prior. `init_to_median` makes more sense then. In Pyro autoguide, we used `init_to_median` by default. However, there are a couple of reasons to stick with `init_to_uniform`
+ `init_to_uniform` is faster (this is a minor point I guess)
+ `init_to_uniform` matches the Stan behavior, which is a desire of many users. We can change this but need to do many experiments by ourselves. :(

I think we should add a clear note the MCMC docs that in some cases, the default init strategy is not good. It is more reliable to use `init_to_median`, `init_to_sample` (for diversity)...

Resolving patch:
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -674,11 +674,11 @@ class seed(Messenger):
 
     def __init__(self, fn=None, rng_seed=None):
         if isinstance(rng_seed, int) or (
-            isinstance(rng_seed, jnp.ndarray) and not jnp.shape(rng_seed)
+            isinstance(rng_seed, (np.ndarray, jnp.ndarray)) and not jnp.shape(rng_seed)
         ):
             rng_seed = random.PRNGKey(rng_seed)
         if not (
-            isinstance(rng_seed, jnp.ndarray)
+            isinstance(rng_seed, (np.ndarray, jnp.ndarray))
             and rng_seed.dtype == jnp.uint32
             and rng_seed.shape == (2,)
         ):
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -2,10 +2,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections import OrderedDict, namedtuple
+from functools import partial
 import math
 import os
 
-from jax import device_put, lax, partial, random, vmap
+from jax import device_put, lax, random, vmap
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 
@@ -511,7 +512,11 @@ class HMC(MCMCKernel):
     Hamiltonian Monte Carlo inference, using fixed trajectory length, with
     provision for step size and mass matrix adaptation.
 
-    .. note:: until the kernel is used in an MCMC run, `postprocess_fn` will return the identity
+    .. note:: Until the kernel is used in an MCMC run, `postprocess_fn` will return the
+        identity function.
+
+    .. note:: The default init strategy ``init_to_uniform`` might not be a good strategy
+        for some models. You might want to try other init strategies like ``init_to_median``.
 
     **References:**
 
@@ -768,7 +773,11 @@ class NUTS(HMC):
     Hamiltonian Monte Carlo inference, using the No U-Turn Sampler (NUTS)
     with adaptive path length and mass matrix adaptation.
 
-    .. note:: until the kernel is used in an MCMC run, `postprocess_fn` will return the identity
+    .. note:: Until the kernel is used in an MCMC run, `postprocess_fn` will return the
+        identity function.
+
+    .. note:: The default init strategy ``init_to_uniform`` might not be a good strategy
+        for some models. You might want to try other init strategies like ``init_to_median``.
 
     **References:**
 
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -7,6 +7,8 @@
 import os
 import warnings
 
+import numpy as np
+
 from jax import jit, lax, local_device_count, pmap, random, vmap
 from jax.core import Tracer
 from jax.interpreters.xla import DeviceArray
@@ -190,7 +192,7 @@ def _hashable(x):
         return x
     elif isinstance(x, DeviceArray):
         return x.copy().tobytes()
-    elif isinstance(x, jnp.ndarray):
+    elif isinstance(x, (np.ndarray, jnp.ndarray)):
         return x.tobytes()
     return x
 
@@ -207,7 +209,8 @@ class MCMC(object):
     .. note:: If setting `num_chains` greater than `1` in a Jupyter Notebook, then you will need to
         have installed `ipywidgets <https://ipywidgets.readthedocs.io/en/latest/user_install.html>`_
         in the environment from which you launced Jupyter in order for the progress bars to render
-        correctly.
+        correctly. If you are using Jupyter Notebook or Jupyter Lab, please also install the
+        corresponding extension package like `widgetsnbextension` or `jupyterlab_widgets`.
 
     :param MCMCKernel sampler: an instance of :class:`~numpyro.infer.mcmc.MCMCKernel` that
         determines the sampler for running MCMC. Currently, only :class:`~numpyro.infer.hmc.HMC`
@@ -234,6 +237,8 @@ class MCMC(object):
     :param bool jit_model_args: If set to `True`, this will compile the potential energy
         computation as a function of model arguments. As such, calling `MCMC.run` again
         on a same sized but different dataset will not result in additional compilation cost.
+        Note that currently, this does not take effect for the case ``num_chains > 1``
+        and ``chain_method == 'parallel'``.
     """
 
     def __init__(


Test patch:
diff --git a/test/contrib/test_optim.py b/test/contrib/test_optim.py
--- a/test/contrib/test_optim.py
+++ b/test/contrib/test_optim.py
@@ -1,10 +1,12 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from functools import partial
+
 from numpy.testing import assert_allclose
 import pytest
 
-from jax import grad, jit, partial, random
+from jax import grad, jit, random
 from jax.lax import fori_loop
 import jax.numpy as jnp
 from jax.test_util import check_close
diff --git a/test/infer/test_compute_downstream_costs.py b/test/infer/test_compute_downstream_costs.py
--- a/test/infer/test_compute_downstream_costs.py
+++ b/test/infer/test_compute_downstream_costs.py
@@ -304,7 +304,7 @@ def test_compute_downstream_costs_big_model_guide_pair(
 
     for k in dc:
         assert guide_trace[k]["log_prob"].shape == dc[k].shape
-        assert_allclose(dc[k], dc_brute[k])
+        assert_allclose(dc[k], dc_brute[k], rtol=2e-7)
 
 
 def plate_reuse_model_guide(include_obs=True, dim1=3, dim2=2):
diff --git a/test/test_optimizers.py b/test/test_optimizers.py
--- a/test/test_optimizers.py
+++ b/test/test_optimizers.py
@@ -1,9 +1,11 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+from functools import partial
+
 import pytest
 
-from jax import grad, jit, partial
+from jax import grad, jit
 import jax.numpy as jnp
 
 from numpyro import optim