Add Dirichlet Multinomial to PyTorch Distributions

[RylanSchaeffer]

🚀 Feature

The Dirichlet-Multinomial distribution is something I need for my research. TensorFlow has this distribution implemented (https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/DirichletMultinomial). Can someone please implement this distribution in PyTorch, or point me towards how I might implement it myself?

Motivation

Common probabilistic language models (e.g. mixtures of unigrams) require using Dirichlet-Multinomial distributions.

Pitch

Add the distribution to Pytorch. That is, there should be a distribution torch.distributions.dirichlet.DirichletMultinomial.

cc @fritzo @neerajprad @alicanb @vishwakftw @nikitaved

[fritzo]

@RylanSchaeffer we have a PyTorch implementation in Pyro, including a sparse optimization (code, docs)

from pyro.distributions import DirichletMultinomial

You can use this directly or move it upstream to PyTorch.

[RylanSchaeffer]

@fritzo , thanks! That might suffice. As a principle, what determines whether something makes it from Pyro to PyTorch?

[fritzo]

what determines whether something makes it from Pyro to PyTorch?

Informally we have moved distributions from Pyro to PyTorch if:

1. there appears to be public demand for the distribution,
2. the distribution's interface is stable,
3. the distribution's implementation is uncontroversial (e.g. does not use unsafe hacks),
4. someone has the motivation and time to migrate code, update docs, and write tests.

I think DirichletMultinomial already satisfies 1-3.