Repository: pyro-ppl/numpyro

Issue:
Incompatible shapes error after stack error
Related to https://github.com/pyro-ppl/numpyro/issues/790

At least under the circumstances listed in #790, when the AssertionError is raised, subsequent inference `run` on the model will fail in an ipython environment. Memoized (compiled?) tensor shapes appears to be corrupted and maintained until the next invocation.

Stack trace from 790:

<details>

```
---------------------------------------------------------------------------
FilteredStackTrace                        Traceback (most recent call last)
<ipython-input-4-6dacd0827004> in <module>
      2 mcmcPP = MCMC(NUTS(model), 100, 1000)
----> 3 mcmcPP.run(init_rng_key, altCounts)
      4 mcmcPP.print_summary()

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py in run(self, rng_key, extra_fields, init_params, *args, **kwargs)
    448         if self.num_chains == 1:
--> 449             states_flat, last_state = partial_map_fn(map_args)
    450             states = tree_map(lambda x: x[jnp.newaxis, ...], states_flat)

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py in _single_chain_mcmc(self, init, args, kwargs, collect_fields)
    312             init_state = self.sampler.init(rng_key, self.num_warmup, init_params,
--> 313                                            model_args=args, model_kwargs=kwargs)
    314         if self.postprocess_fn is None:

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py in init(self, rng_key, num_warmup, init_params, model_args, model_kwargs)
    449             rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1)
--> 450         init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params)
    451         if self._potential_fn and init_params is None:

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py in _init_state(self, rng_key, model_args, model_kwargs, init_params)
    412                 model_args=model_args,
--> 413                 model_kwargs=model_kwargs)
    414             if self._init_fn is None:

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in initialize_model(rng_key, model, init_strategy, dynamic_args, model_args, model_kwargs)
    429                                                                   model_kwargs=model_kwargs,
--> 430                                                                   prototype_params=prototype_params)
    431 

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in find_valid_initial_params(rng_key, model, init_strategy, enum, model_args, model_kwargs, prototype_params)
    256     if rng_key.ndim == 1:
--> 257         (init_params, pe, z_grad), is_valid = _find_valid_params(rng_key, exit_early=True)
    258     else:

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in _find_valid_params(rng_key, exit_early)
    244             # where we can avoid compiling body_fn in while_loop.
--> 245             _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
    246             if not_jax_tracer(is_valid):

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in body_fn(state)
    234         potential_fn = partial(potential_energy, model, model_args, model_kwargs, enum=enum)
--> 235         pe, z_grad = value_and_grad(potential_fn)(params)
    236         z_grad_flat = ravel_pytree(z_grad)[0]

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in potential_energy(model, model_args, model_kwargs, params, enum)
    156     # no param is needed for log_density computation because we already substitute
--> 157     log_joint, model_trace = log_density_(substituted_model, model_args, model_kwargs, {})
    158     return - log_joint

~/miniconda3/lib/python3.7/site-packages/numpyro/contrib/funsor/infer_util.py in log_density(model, model_args, model_kwargs, params)
    123     with plate_to_enum_plate():
--> 124         model_trace = packed_trace(model).get_trace(*model_args, **model_kwargs)
    125     log_factors = []

~/miniconda3/lib/python3.7/site-packages/numpyro/handlers.py in get_trace(self, *args, **kwargs)
    157         """
--> 158         self(*args, **kwargs)
    159         return self.trace

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __call__(self, *args, **kwargs)
     67         with self:
---> 68             return self.fn(*args, **kwargs)
     69 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __exit__(self, *args, **kwargs)
     56     def __exit__(self, *args, **kwargs):
---> 57         assert _PYRO_STACK[-1] is self
     58         _PYRO_STACK.pop()

FilteredStackTrace: AssertionError

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exceptio (where 6,4 is the dimension of the pooled 'conc' sample statement, and 6 is K):

AssertionError                            Traceback (most recent call last)
<ipython-input-4-6dacd0827004> in <module>
      1 init_rng_key = random.PRNGKey(12273)
      2 mcmcPP = MCMC(NUTS(model), 100, 1000)
----> 3 mcmcPP.run(init_rng_key, altCounts)
      4 mcmcPP.print_summary()

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py in run(self, rng_key, extra_fields, init_params, *args, **kwargs)
    447         map_args = (rng_key, init_state, init_params)
    448         if self.num_chains == 1:
--> 449             states_flat, last_state = partial_map_fn(map_args)
    450             states = tree_map(lambda x: x[jnp.newaxis, ...], states_flat)
    451         else:

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py in _single_chain_mcmc(self, init, args, kwargs, collect_fields)
    311         if init_state is None:
    312             init_state = self.sampler.init(rng_key, self.num_warmup, init_params,
--> 313                                            model_args=args, model_kwargs=kwargs)
    314         if self.postprocess_fn is None:
    315             postprocess_fn = self.sampler.postprocess_fn(args, kwargs)

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py in init(self, rng_key, num_warmup, init_params, model_args, model_kwargs)
    448         else:
    449             rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1)
--> 450         init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params)
    451         if self._potential_fn and init_params is None:
    452             raise ValueError('Valid value of `init_params` must be provided with'

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py in _init_state(self, rng_key, model_args, model_kwargs, init_params)
    411                 init_strategy=self._init_strategy,
    412                 model_args=model_args,
--> 413                 model_kwargs=model_kwargs)
    414             if self._init_fn is None:
    415                 self._init_fn, self._sample_fn = hmc(potential_fn_gen=potential_fn,

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in initialize_model(rng_key, model, init_strategy, dynamic_args, model_args, model_kwargs)
    428                                                                   model_args=model_args,
    429                                                                   model_kwargs=model_kwargs,
--> 430                                                                   prototype_params=prototype_params)
    431 
    432     if not_jax_tracer(is_valid):

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in find_valid_initial_params(rng_key, model, init_strategy, enum, model_args, model_kwargs, prototype_params)
    255     # Handle possible vectorization
    256     if rng_key.ndim == 1:
--> 257         (init_params, pe, z_grad), is_valid = _find_valid_params(rng_key, exit_early=True)
    258     else:
    259         (init_params, pe, z_grad), is_valid = lax.map(_find_valid_params, rng_key)

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in _find_valid_params(rng_key, exit_early)
    243             # Early return if valid params found. This is only helpful for single chain,
    244             # where we can avoid compiling body_fn in while_loop.
--> 245             _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
    246             if not_jax_tracer(is_valid):
    247                 if device_get(is_valid):

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in body_fn(state)
    233 
    234         potential_fn = partial(potential_energy, model, model_args, model_kwargs, enum=enum)
--> 235         pe, z_grad = value_and_grad(potential_fn)(params)
    236         z_grad_flat = ravel_pytree(z_grad)[0]
    237         is_valid = jnp.isfinite(pe) & jnp.all(jnp.isfinite(z_grad_flat))

~/miniconda3/lib/python3.7/site-packages/jax/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    137   def reraise_with_filtered_traceback(*args, **kwargs):
    138     try:
--> 139       return fun(*args, **kwargs)
    140     except Exception as e:
    141       if not is_under_reraiser(e):

~/miniconda3/lib/python3.7/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    808     tree_map(partial(_check_input_dtype_grad, holomorphic, allow_int), dyn_args)
    809     if not has_aux:
--> 810       ans, vjp_py = _vjp(f_partial, *dyn_args)
    811     else:
    812       ans, vjp_py, aux = _vjp(f_partial, *dyn_args, has_aux=True)

~/miniconda3/lib/python3.7/site-packages/jax/api.py in _vjp(fun, has_aux, *primals)
   1839   if not has_aux:
   1840     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 1841     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1842     out_tree = out_tree()
   1843   else:

~/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    111 def vjp(traceable, primals, has_aux=False):
    112   if not has_aux:
--> 113     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    114   else:
    115     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

~/miniconda3/lib/python3.7/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     98   _, in_tree = tree_flatten(((primals, primals), {}))
     99   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
--> 100   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
    101   out_primals_pvals, out_tangents_pvals = tree_unflatten(out_tree(), out_pvals)
    102   assert all(out_primal_pval.is_known() for out_primal_pval in out_primals_pvals)

~/miniconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate)
    402   with core.new_main(JaxprTrace) as main:
    403     fun = trace_to_subjaxpr(fun, main, instantiate)
--> 404     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    405     assert not env
    406     del main

~/miniconda3/lib/python3.7/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    154 
    155     try:
--> 156       ans = self.f(*args, **dict(self.params, **kwargs))
    157     except:
    158       # Some transformations yield from inside context managers, so we have to

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in potential_energy(model, model_args, model_kwargs, params, enum)
    155     substituted_model = substitute(model, substitute_fn=partial(_unconstrain_reparam, params))
    156     # no param is needed for log_density computation because we already substitute
--> 157     log_joint, model_trace = log_density_(substituted_model, model_args, model_kwargs, {})
    158     return - log_joint
    159 

~/miniconda3/lib/python3.7/site-packages/numpyro/contrib/funsor/infer_util.py in log_density(model, model_args, model_kwargs, params)
    122     model = substitute(model, data=params)
    123     with plate_to_enum_plate():
--> 124         model_trace = packed_trace(model).get_trace(*model_args, **model_kwargs)
    125     log_factors = []
    126     time_to_factors = defaultdict(list)  # log prob factors

~/miniconda3/lib/python3.7/site-packages/numpyro/handlers.py in get_trace(self, *args, **kwargs)
    156         :return: `OrderedDict` containing the execution trace.
    157         """
--> 158         self(*args, **kwargs)
    159         return self.trace
    160 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __call__(self, *args, **kwargs)
     66     def __call__(self, *args, **kwargs):
     67         with self:
---> 68             return self.fn(*args, **kwargs)
     69 
     70 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __exit__(self, *args, **kwargs)
     55 
     56     def __exit__(self, *args, **kwargs):
---> 57         assert _PYRO_STACK[-1] is self
     58         _PYRO_STACK.pop()
     59 

AssertionError: 
```

</details>

After fixing the cause of the error (wrapping the pooled sample statement in a plate, or removing the offending plate), a shape error is generated with the following stack trace:

<details>

```
~/miniconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py in run(self, rng_key, extra_fields, init_params, *args, **kwargs)
    447         map_args = (rng_key, init_state, init_params)
    448         if self.num_chains == 1:
--> 449             states_flat, last_state = partial_map_fn(map_args)
    450             states = tree_map(lambda x: x[jnp.newaxis, ...], states_flat)
    451         else:

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py in _single_chain_mcmc(self, init, args, kwargs, collect_fields)
    311         if init_state is None:
    312             init_state = self.sampler.init(rng_key, self.num_warmup, init_params,
--> 313                                            model_args=args, model_kwargs=kwargs)
    314         if self.postprocess_fn is None:
    315             postprocess_fn = self.sampler.postprocess_fn(args, kwargs)

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py in init(self, rng_key, num_warmup, init_params, model_args, model_kwargs)
    448         else:
    449             rng_key, rng_key_init_model = jnp.swapaxes(vmap(random.split)(rng_key), 0, 1)
--> 450         init_params = self._init_state(rng_key_init_model, model_args, model_kwargs, init_params)
    451         if self._potential_fn and init_params is None:
    452             raise ValueError('Valid value of `init_params` must be provided with'

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/hmc.py in _init_state(self, rng_key, model_args, model_kwargs, init_params)
    411                 init_strategy=self._init_strategy,
    412                 model_args=model_args,
--> 413                 model_kwargs=model_kwargs)
    414             if self._init_fn is None:
    415                 self._init_fn, self._sample_fn = hmc(potential_fn_gen=potential_fn,

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in initialize_model(rng_key, model, init_strategy, dynamic_args, model_args, model_kwargs)
    393                                    substitute_fn=init_strategy)
    394     inv_transforms, replay_model, has_enumerate_support, model_trace = _get_model_transforms(
--> 395         substituted_model, model_args, model_kwargs)
    396     # substitute param sites from model_trace to model so
    397     # we don't need to generate again parameters of `numpyro.module`

~/miniconda3/lib/python3.7/site-packages/numpyro/infer/util.py in _get_model_transforms(model, model_args, model_kwargs)
    263 def _get_model_transforms(model, model_args=(), model_kwargs=None):
    264     model_kwargs = {} if model_kwargs is None else model_kwargs
--> 265     model_trace = trace(model).get_trace(*model_args, **model_kwargs)
    266     inv_transforms = {}
    267     # model code may need to be replayed in the presence of deterministic sites

~/miniconda3/lib/python3.7/site-packages/numpyro/handlers.py in get_trace(self, *args, **kwargs)
    156         :return: `OrderedDict` containing the execution trace.
    157         """
--> 158         self(*args, **kwargs)
    159         return self.trace
    160 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __call__(self, *args, **kwargs)
     66     def __call__(self, *args, **kwargs):
     67         with self:
---> 68             return self.fn(*args, **kwargs)
     69 
     70 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __call__(self, *args, **kwargs)
     66     def __call__(self, *args, **kwargs):
     67         with self:
---> 68             return self.fn(*args, **kwargs)
     69 
     70 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in __call__(self, *args, **kwargs)
     66     def __call__(self, *args, **kwargs):
     67         with self:
---> 68             return self.fn(*args, **kwargs)
     69 
     70 

<ipython-input-5-2f909eaf93fa> in modelPartialPooled(data)
    177 def modelPartialPooled(data):
    178     # This also works, for a single set of shraed parameters
--> 179     conc = numpyro.sample("conc", Gamma(pdsAllShaped, 1))
    180     # print("conc.shape", conc)
    181     with numpyro.plate("beta_plate", nHypotheses):

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in sample(name, fn, obs, rng_key, sample_shape, infer)
    112 
    113     # ...and use apply_stack to send it to the Messengers
--> 114     msg = apply_stack(initial_msg)
    115     return msg['value']
    116 

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in apply_stack(msg)
     22     pointer = 0
     23     for pointer, handler in enumerate(reversed(_PYRO_STACK)):
---> 24         handler.process_message(msg)
     25         # When a Messenger sets the "stop" field of a message,
     26         # it prevents any Messengers above it on the stack from being applied.

~/miniconda3/lib/python3.7/site-packages/numpyro/contrib/funsor/enum_messenger.py in process_message(self, msg)
    470         if msg["type"] in ["to_funsor", "to_data"]:
    471             return super().process_message(msg)
--> 472         return OrigPlateMessenger.process_message(self, msg)
    473 
    474     def postprocess_message(self, msg):

~/miniconda3/lib/python3.7/site-packages/numpyro/primitives.py in process_message(self, msg)
    314             overlap_idx = max(len(expected_shape) - len(dist_batch_shape), 0)
    315             trailing_shape = expected_shape[overlap_idx:]
--> 316             broadcast_shape = lax.broadcast_shapes(trailing_shape, tuple(dist_batch_shape))
    317             batch_shape = expected_shape[:overlap_idx] + broadcast_shape
    318             msg['fn'] = msg['fn'].expand(batch_shape)

~/miniconda3/lib/python3.7/site-packages/jax/lax/lax.py in broadcast_shapes(*shapes)
     80   if result_shape is None:
     81     raise ValueError("Incompatible shapes for broadcasting: {}"
---> 82                      .format(tuple(map(tuple, shapes))))
     83   return result_shape
     84 

ValueError: Incompatible shapes for broadcasting: ((1, 6), (6, 4))
```

</details>

Reproduction case:

```python
def modelPartialPooled(data):
    # to fix get rid of beta_plate and z
    conc = numpyro.sample("conc", Gamma(1, 1))
    with numpyro.plate("beta_plate", K):
        beta = numpyro.sample("beta", Beta(1, alpha, validate_args=False))

    with numpyro.plate("data", N):
        z = numpyro.sample("z", Categorical(mix_weights(beta)))
        # z specifically not used, to reduce degrees of freedom of the test
        probs = numpyro.sample("probs", Dirichlet(conc[0], validate_args=False))
        return numpyro.sample("obs", Multinomial(probs=probs, validate_args=False), obs=data)
```




Discussion:
It seems that if a discrete variable is not used elsewhere, then the error will happen. We should catch that issue and raise a better error message. Thanks, @akotlar!

cc @eb8680
> [after] the `AssertionError` is raised, subsequent inference ...

I believe the `_PYRO_STACK` is not being correctly cleaned on error, due to logic in the `__exit__()` methods of handlers.
Thanks, @fritzo! I understand the situtation now. Will make the fix soon.

Resolving patch:
diff --git a/numpyro/contrib/funsor/infer_util.py b/numpyro/contrib/funsor/infer_util.py
--- a/numpyro/contrib/funsor/infer_util.py
+++ b/numpyro/contrib/funsor/infer_util.py
@@ -175,4 +175,8 @@ def model(*args, **kwargs):
             funsor.ops.logaddexp, funsor.ops.add, log_factors,
             eliminate=sum_vars | prod_vars, plates=prod_vars)
     result = funsor.optimizer.apply_optimizer(lazy_result)
+    if len(result.inputs) > 0:
+        raise ValueError("Expected the joint log density is a scalar, but got {}. "
+                         "There seems to be something wrong at the following sites: {}."
+                         .format(result.data.shape, {k.split("__BOUND")[0] for k in result.inputs}))
     return result.data, model_trace
diff --git a/numpyro/handlers.py b/numpyro/handlers.py
--- a/numpyro/handlers.py
+++ b/numpyro/handlers.py
@@ -278,18 +278,23 @@ def __enter__(self):
         COERCIONS.append(self._coerce)
         return super().__enter__()
 
-    def __exit__(self, *args, **kwargs):
+    def __exit__(self, exc_type, exc_value, traceback):
         import funsor
 
         _coerce = COERCIONS.pop()
         assert _coerce is self._coerce
-        super().__exit__(*args, **kwargs)
+        super().__exit__(exc_type, exc_value, traceback)
+
+        if exc_type is not None:
+            return
 
         # Convert delayed statements to pyro.factor()
         reduced_vars = []
         log_prob_terms = []
         plates = frozenset()
         for name, site in self.trace.items():
+            if site["type"] != "sample":
+                continue
             if not site["is_observed"]:
                 reduced_vars.append(name)
             dim_to_name = {f.dim: f.name for f in site["cond_indep_stack"]}
diff --git a/numpyro/primitives.py b/numpyro/primitives.py
--- a/numpyro/primitives.py
+++ b/numpyro/primitives.py
@@ -52,9 +52,21 @@ def __init__(self, fn=None):
     def __enter__(self):
         _PYRO_STACK.append(self)
 
-    def __exit__(self, *args, **kwargs):
-        assert _PYRO_STACK[-1] is self
-        _PYRO_STACK.pop()
+    def __exit__(self, exc_type, exc_value, traceback):
+        if exc_type is None:
+            assert _PYRO_STACK[-1] is self
+            _PYRO_STACK.pop()
+        else:
+            # NB: this mimics Pyro exception handling
+            # the wrapped function or block raised an exception
+            # handler exception handling:
+            # when the callee or enclosed block raises an exception,
+            # find this handler's position in the stack,
+            # then remove it and everything below it in the stack.
+            if self in _PYRO_STACK:
+                loc = _PYRO_STACK.index(self)
+                for i in range(loc, len(_PYRO_STACK)):
+                    _PYRO_STACK.pop()
 
     def process_message(self, msg):
         pass


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
@@ -21,6 +21,7 @@
 from numpyro.contrib.indexing import Vindex
 import numpyro.distributions as dist
 from numpyro.infer import MCMC, NUTS
+from numpyro.primitives import _PYRO_STACK
 
 
 def test_gaussian_mixture_model():
@@ -432,7 +433,7 @@ def transition_fn(name, probs, locs, x, y):
     assert_allclose(actual_log_joint, expected_log_joint)
 
 
-def test_missing_plate():
+def test_missing_plate(monkeypatch):
     K, N = 3, 1000
 
     def gmm(data):
@@ -453,3 +454,8 @@ def gmm(data):
     mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=500)
     with pytest.raises(AssertionError, match="Missing plate statement"):
         mcmc.run(random.PRNGKey(2), data)
+
+    monkeypatch.setattr(numpyro.infer.util, "_validate_model", lambda model_trace: None)
+    with pytest.raises(Exception):
+        mcmc.run(random.PRNGKey(2), data)
+    assert len(_PYRO_STACK) == 0