Add models for distributions

Have dataclasses for torch.distributions . This would be helpful especially for any generative modelling where you might need to sample from different distributions to check which work best as the latent input.