Add format_shapes utility

[tcbegley]

This is a proposed solution to #1104. It's a pretty direct port of .format_shapes from Pyro.

I've created a gist with examples from the Pyro docs.

Would love some feedback on a few things in particular:

• Function name + location. I just put it in numpyro.util for now.
• The main difference between this and Pyro that I'm aware of is that in Pyro plates are traced as "sample" sites, which means they show up in the result of format_shapes. I can add some logic to add "plate" sites from NumPyro traces to the result. Should they be mixed in with the sample sites like in Pyro or should they get their own section?
• Added a log_prob argument to mimic the effect of running trace.compute_log_prob() before trace.format_shapes() in Pyro. Is this a reasonable solution?
• Couldn't see any tests in the Pyro repo for this to replicate. Is it worth writing tests for?

TODO:

• Update docstring

[fehiepsi]

Is it worth writing tests for?

I think you can compare the output with expected strings

Should they be mixed in with the sample sites like in Pyro

I think this is the way to go and I like all other resolutions. :)

[fehiepsi]

Looks great to me! Thank you, this utility is really helpful!