Repository: pyro-ppl/numpyro

Issue:
[ENH, DOCS] Add tutorial on how to deal with discrete latent variables
Numpyro is great! But it took me a while to figure out how to work with discrete latent variables. Unfortunately, there's not much in the documentation about it yet.

It would probably be good to condense the various discussions in issues, gists and forum posts into a doc page.

These are the resources that helped me - if anyone stumbles upon this issue, here you go ;)
Arviz workaround: #1121 
How to get discrete samples: https://github.com/pyro-ppl/numpyro/issues/770
Tutorial: https://gist.github.com/peterroelants/277d9d47a76a55c23a433d37bbbd6dd7/732b2af0881f96a0175a5f79ca556f3a48996e7a
Example in doc (not dedicated to discrete latent variable): https://num.pyro.ai/en/stable/examples/annotation.html


Discussion:
> not dedicated to discrete latent variable

I guess your request is to have a tutorial that works with `arviz`? Currently, we have [5 examples/tutorials](https://num.pyro.ai/en/stable/) for models with discrete latent variables that cover various scenario: mixture model, hidden markov model, time series, imputation, non-enumerate (with discrete hmc gibbs in the nested sampling example). I think we can add some codes to the `annotation` example to illustrate that usage case (I guess it might be better to have such example/tutorial in arviz but not many frameworks support inference with discrete latent variables through marginalization). Do you want to contribute? :)
100% agree that numpyro is great!  I've been using pymc2 long after it has been deprecated, and I'm think about switching to numpyro.

If you do write a documentation on this, can you help me understand why plates are needed?  I.e. why doesn't something like this version work?

```
def gmm_model_no_plate(data, k):
    selection_prob = numpyro.sample('selection_prob', dist.Dirichlet(concentration=jnp.ones(k)))
    
    mu = numpyro.sample('mu', dist.Normal(loc=jnp.zeros(k), scale=10.*jnp.ones(k)))
    sigma = numpyro.sample('scale', dist.HalfCauchy(scale=10*jnp.ones(k)))
    
    cluster_idx = numpyro.sample('cluster_idx', dist.Categorical(selection_prob))
    numpyro.sample('x', dist.Normal(loc=mu[cluster_idx], scale=sigma[cluster_idx]), obs=data)
```


Hi @aflaxman, that version will work with algorithms like [DiscreteHMCGibbs](https://num.pyro.ai/en/stable/mcmc.html#numpyro.infer.hmc_gibbs.DiscreteHMCGibbs) or `MixedHMC`. But for algorithms that require [enumeration](https://pyro.ai/examples/enumeration.html), at each site we will need to specify which dimension we need to marginalize over. For example, for each value of `cluster_idx`, we have a corresponding likelihood of `x`. Under enumeration, `x` will have an additional `cluster_idx`, and we will need to specify which dimension of `x` that we want to marginalize (i.e. taking integration).

`plate` is useful for many other usage cases, to name a few:
+ it clearly denotes the independent relationship of variables
+ it helps users to avoid many shape errors, which FWIK is the most common error when creating a probabilistic model, in particularly
    + it makes a model easier to [visualize](https://num.pyro.ai/en/stable/tutorials/model_rendering.html) and [inspect](https://num.pyro.ai/en/stable/utilities.html#numpyro.util.format_shapes)
+ it helps the code easier to read, in particularly
    + it simplifies many broadcasting computations (e.g., we can use `dist.Normal(0, 10)` rather than `dist.Normal(jnp.zeros(k), 10 * jnp.ones(k))`)
    + the code should correspond 1-1 to [plate notation](https://en.wikipedia.org/wiki/Plate_notation)

Resolving patch:
diff --git a/examples/annotation.py b/examples/annotation.py
--- a/examples/annotation.py
+++ b/examples/annotation.py
@@ -332,6 +332,24 @@ def main(args):
         print(row_format.format(f"item[{i}]", *row))
 
 
+# %%
+# .. note::
+#     In the above inference code, we marginalized the discrete latent variables `c`
+#     hence `mcmc.get_samples(...)` does not include samples of `c`. We then utilize
+#     `Predictive(..., infer_discrete=True)` to get posterior samples for `c`, which
+#     is stored in `discrete_samples`. To merge those discrete samples into the `mcmc`
+#     instance, we can use the following pattern::
+#
+#         chain_discrete_samples = jax.tree_util.tree_map(
+#             lambda x: x.reshape((args.num_chains, args.num_samples) + x.shape[1:]),
+#             discrete_samples)
+#         mcmc.get_samples().update(discrete_samples)
+#         mcmc.get_samples(group_by_chain=True).update(chain_discrete_samples)
+#
+#     This is useful when we want to pass the `mcmc` instance to `arviz` through
+#     `arviz.from_numpyro(mcmc)`.
+
+
 if __name__ == "__main__":
     assert numpyro.__version__.startswith("0.8.0")
     parser = argparse.ArgumentParser(description="Bayesian Models of Annotation")


Test patch: