Expose energy for bfmi diagnostic

[fehiepsi]

Resolves #349.

This PR exposes energy, which is required for bfmi diagnostic, and energy plot. Because this is just a scalar, I think that storing this value to TreeInfo won't affect the performance.

@neerajprad This gist illustrates how numpyro and arviz are integrated. With this PR, we fully support arviz (I need to make a PR to arviz with those updates after 0.2.1 is released)! Cheers, :D

I think it is good to make a tutorial which shows how simple it is to run MCMC in numpyro, and showcase the support for arviz. WDYT?

[neerajprad]

@neerajprad This gist illustrates how numpyro and arviz are integrated. With this PR, we fully support arviz (I need to make a PR to arviz with those updates after 0.2.1 is released)! Cheers, :D

That's great! @null-a - I think we should have these utilities available within the next release for numpyro and arviz.

I think it is good to make a tutorial which shows how simple it is to run MCMC in numpyro, and showcase the support for arviz. WDYT?

Sure! We can modify the existing examples too. I think one good example is the existing one on non-centered parametrization which can be converted into a tutorial that demonstrates this integration and shows how to use diagnostic information to detect issues with the model.