Add support for named dims (torchdim)

[ordabayevy]

Few notes about this PR

• Importing functorch.dim messes up torch.Tensor a bit (e.g., the torch.Tensor.split method starts to fail). Therefore, I siloed the functorch.dim import from the main Pyro
  • That's why I also removed pyro/contrib/named from doctest in order to avoid the import of functorch.dim
  • And use hasattr(self.dim, "is_bound") instead of isinstance(self.dim, Dim)
• Trace_ELBO implementation is actually similar to TraceGraph_ELBO with dependency tracking with the goal of having a single Trace_ELBO implementation that will generalize to TraceEnum_ELBO and others.
• pyro.enable_validation needs to be set to False to avoid validation errors caused by if value.all() method
• Distribution arguments all need to be bound by named dim, otherwise broadcasting attempt of parameters by a distribution will lead to segmentation fault

[ordabayevy]

Unit tests for distribution shapes for log_prob, mean, sample, rsample, entropy (fail when named and positional dims are mixed in the batch/event/sample shape; conflicting named dims)

Generalize named dim binding implementation.

Test transforms and support.

Shape inference.

[ordabayevy]

Add type annotation. Move to contrib/named in the follow up PR.