Support exporting to arviz InferenceData about numpyro

[fehiepsi]

Actually, our fori_append output + transform functionalities are quite convenient for this job. The following code plots well the result.

hmc_states = fori_append(sample_kernel, hmc_state, num_samples,
                         transform=lambda x: {"z": x.z})

import arviz as az
data = arviz.from_dict(hmc_states)
az.plot_trace(data);

@neerajprad I think we should make a notebook for bayesian regression using arviz instead of making custom functions as in pyro's bayesian regression :)

[neerajprad]

Good idea. I think since we are using very straightforward data structures, it is best to delegate all the diagnostics / plotting to arviz.