Repository: pyro-ppl/numpyro

Issue:
Initialization issues when using mask/latent variables
### Description
I am not sure where the problem is coming from exactly, either from discrete latent variables or using `handlers.mask`. A relatively small script that reproduces it is below: I try to infer in a mixture model, where the mixture ID is determined by a Bernoulli variable. I use `mask` to decide between sampling from either of the two distributions.
When I run MCMC inference with `init_strategy=numpyro.infer.init_to_uniform`, for both models I get `xyz.shape == (7, 100, 3)` (as it should be). When I change to `init_to_median`, for `model2` I get `xyz.shape == (7, 2, 100, 3)`. In the full version (before condensing it to make it self-contained) the behavior was even more esoteric, with the extra dimension appearing in `model2` when I specify any `init_strategy`, including `init_to_uniform`, which is the default!

### Reproduction

```python
import numpy as np
import jax
import jax.numpy as jnp
import numpyro as npyro
import numpyro.distributions as dist

def f(xyz):
    ''' an arbitrary function from R^3 -> R^2'''
    return jnp.cos(xyz)[..., :2]

def g(xyz):
    ''' another arbitrary function from R^3 -> R^2'''
    return jnp.sin(xyz)[..., :2]

def model1(e1, e2):
    ''' a single component baseline '''
    n = e1.shape[0]

    obs_δ = npyro.sample('obs_δ', dist.LogNormal(scale=1.))
    obs_cov = jnp.eye(2) * obs_δ

    with npyro.plate('n', n):
        xyz = npyro.sample(
            'xyz',
            dist.MultivariateNormal(jnp.ones(3), jnp.eye(3)),
        )

        obs_1 = npyro.sample(
            'obs_1',
            dist.MultivariateNormal(
                loc=f(xyz),
                covariance_matrix=obs_cov,
            ),
            obs=e1,
        )
        obs_2 = npyro.sample(
            'obs_2',
            dist.MultivariateNormal(
                loc=g(xyz),
                covariance_matrix=obs_cov,
            ),
            obs=e2,
        )

def model2(e1, e2):
    ''' a two component mixture '''
    n = e1.shape[0]

    obs_δ = npyro.sample('obs_δ', dist.LogNormal(scale=1.))
    obs_cov = jnp.eye(2) * obs_δ

    cmp_p = npyro.sample('cmp_p', dist.Beta(1, 1))

    with npyro.plate('n', n):
        is_cmp_1 = npyro.sample(
            'is_cmp_1',
            dist.Bernoulli(cmp_p),
        ).astype(bool)

        with npyro.handlers.mask(mask=is_cmp_1):
            xyz = npyro.sample(
                'xyz',
                dist.MultivariateNormal(jnp.ones(3), jnp.eye(3)),
            )

            obs_1 = npyro.sample(
                'obs_1_1',
                dist.MultivariateNormal(
                    loc=f(xyz),
                    covariance_matrix=obs_cov,
                ),
                obs=e1,
            )
            obs_2 = npyro.sample(
                'obs_1_2',
                dist.MultivariateNormal(
                    loc=g(xyz),
                    covariance_matrix=obs_cov,
                ),
                obs=e2,
            )

        with npyro.handlers.mask(mask=~is_cmp_1):
            obs_1 = npyro.sample(
                'obs_2_1',
                dist.Uniform(
                    low=jnp.zeros(2),
                    high=jnp.ones(2),
                ).to_event(1),
                obs=e1,
            )
            obs_2 = npyro.sample(
                'obs_2_2',
                dist.Uniform(
                    low=jnp.zeros(2),
                    high=jnp.ones(2),
                ).to_event(1),
                obs=e2,
            )

def run_mcmc(model, *evidence):
    rng_key = jax.random.PRNGKey(42)
    kernel = npyro.infer.NUTS(model, init_strategy=npyro.infer.init_to_median)
    #kernel = npyro.infer.NUTS(model) # this would give the same output shapes for both models
    mcmc = npyro.infer.MCMC(
        kernel,
        num_warmup=2,
        num_samples=7,
        num_chains=1,
        progress_bar=True,
    )
    mcmc.run(rng_key, *evidence)
    predictions = mcmc.get_samples()

    print(20 * '-' + ' ' + model.__name__ + ' ' + 20 * '-')
    for key, value in predictions.items():
        print(key, value.shape)
    
evidence = (np.random.randn(100, 2), np.random.randn(100, 2))

run_mcmc(model1, *evidence)
run_mcmc(model2, *evidence)
```

### Setup
I am running numpyro 0.6.0 with jax 0.2.10.


Discussion:

Resolving patch:
diff --git a/numpyro/infer/util.py b/numpyro/infer/util.py
--- a/numpyro/infer/util.py
+++ b/numpyro/infer/util.py
@@ -239,8 +239,18 @@ def body_fn(state):
         key, subkey = random.split(key)
 
         if radius is None or prototype_params is None:
+            # XXX: we don't want to apply enum to draw latent samples
+            model_ = model
+            if enum:
+                from numpyro.contrib.funsor import enum as enum_handler
+
+                if isinstance(model, substitute) and isinstance(model.fn, enum_handler):
+                    model_ = substitute(model.fn.fn, data=model.data)
+                elif isinstance(model, enum_handler):
+                    model_ = model.fn
+
             # Wrap model in a `substitute` handler to initialize from `init_loc_fn`.
-            seeded_model = substitute(seed(model, subkey), substitute_fn=init_strategy)
+            seeded_model = substitute(seed(model_, subkey), substitute_fn=init_strategy)
             model_trace = trace(seeded_model).get_trace(*model_args, **model_kwargs)
             constrained_values, inv_transforms = {}, {}
             for k, v in model_trace.items():


Test patch:
diff --git a/test/contrib/test_funsor.py b/test/contrib/test_funsor.py
--- a/test/contrib/test_funsor.py
+++ b/test/contrib/test_funsor.py
@@ -19,7 +19,7 @@
 from numpyro.contrib.funsor.infer_util import log_density
 from numpyro.contrib.indexing import Vindex
 import numpyro.distributions as dist
-from numpyro.infer import MCMC, NUTS
+from numpyro.infer import MCMC, NUTS, init_to_median
 from numpyro.primitives import _PYRO_STACK
 
 
@@ -49,7 +49,7 @@ def gmm(data):
         random.PRNGKey(1)
     )
 
-    nuts_kernel = NUTS(gmm)
+    nuts_kernel = NUTS(gmm, init_strategy=init_to_median)
     mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=500)
     mcmc.run(random.PRNGKey(2), data)
     samples = mcmc.get_samples()