text: Title: [FR] add format_shapes method to trace
  Body: Having `trace.format_shapes()` as in [Pyro](https://docs.pyro.ai/en/stable/poutine.html?highlight=format_shapes#pyro.poutine.Trace.format_shapes) would help with debugging and verifying models, especially with enumeration, as shown in the [Pyro tutorial](https://pyro.ai/examples/tensor_shapes.html#Broadcasting-to-allow-parallel-enumeration).