Illustrate the usage of infer_discrete in the annotation example

[fehiepsi]

It seems that we have some infer_discrete bug in the master branch. I merged changes in #991 and it works for models in this example.

Changes in this PR:

• refactor infer_discrete following Make infer_discrete work with scan #991
• deprecate replace handler's guide_trace argument (to be consistent with Pyro)

Questions for reviewers: I found that it is tricky for users to perform infer_discrete correctly to get extra posterior samples for discrete latent sites. Should we make some helpers like MCMC.infer_discrete or Predictive.infer_discrete that does the job in this annotation example under the hood?

    def infer_discrete_model(rng_key, samples):
        conditioned_model = handlers.condition(model, data=samples)
        infer_discrete_model = infer_discrete(
            config_enumerate(conditioned_model), rng_key=rng_key
        )
        with handlers.trace() as tr:
            infer_discrete_model(*data)

        return {
            name: site["value"]
            for name, site in tr.items()
            if site["type"] == "sample" and site["infer"].get("enumerate") == "parallel"
        }

    posterior_samples = mcmc.get_samples()
    discrete_samples = vmap(infer_discrete_model)(
        random.split(random.PRNGKey(1), args.num_samples), posterior_samples
    )

I imagine that it would be convenient if we have something like

mcmc = MCMC(...)
mcmc.run()
mcmc.infer_discrete(rng_key, *args, first_available_dim=None, temperature=1, **kwargs)
posterior_samples = mcmc.get_samples()  # this will include discrete latent sites

[eb8680]

Looks good to me.

Should we make some helpers like MCMC.infer_discrete or Predictive.infer_discrete that does the job in this annotation example under the hood?

Yes - in fact, this should probably be the default behavior of Predictive for HMC/NUTS when infer_discrete is finished.

[eb8680]

Does using lazy here as in the old version and funsor.adjoint.adjoint below not work? I see we made this change in #991 but can't recall why.

[eb8680]

Does using the more functional interface funsor.adjoint.adjoint not work here?

[fehiepsi]

Sorry for the late response, @eb8680! Yes, for some reason there is an error happens at this line.

  File "/home/fehiepsi/numpyro/numpyro/contrib/funsor/discrete.py", line 120, in _sample_posterior
    approx_factors = funsor.adjoint.adjoint(sum_op, prod_op, log_prob)
  File "/home/fehiepsi/funsor/funsor/adjoint.py", line 135, in adjoint
    return tape.adjoint(sum_op, bin_op, root)
  File "/home/fehiepsi/funsor/funsor/adjoint.py", line 115, in adjoint
    in_adjs = adjoint_ops(fn, sum_op, bin_op, adjoint_values[output], *inputs)
  File "/home/fehiepsi/funsor/funsor/registry.py", line 106, in __call__
    return self[key](*args)
  File "/home/fehiepsi/funsor/funsor/registry.py", line 63, in __call__
    return self.partial_call(*args)(*args)
  File "/home/fehiepsi/funsor/funsor/adjoint.py", line 209, in adjoint_contract_generic
    assert len(terms) == 1 or len(terms) == 2
AssertionError

[fehiepsi]

Thanks for helping me review this PR, @eb8680!!