Numpyro: Implement some popular distributions

[fehiepsi]

This issue tracks some popular distributions that are not available in NumPyro yet.

• [x] [Geometric distribution](https://pytorch.org/docs/master/distributions.html#geometric)
• [ ] [ZeroInflatedNegativeBinomial](http://docs.pyro.ai/en/stable/distributions.html#zeroinflatednegativebinomial)
• [x] [Weibull distribution](https://pytorch.org/docs/master/distributions.html#weibull) (supported in contrib.tfp)
• [x] [DirichletMultinomial](http://docs.pyro.ai/en/stable/distributions.html#dirichletmultinomial) (supported in contrib.tfp)
• [ ] [FoldedDistribution](http://docs.pyro.ai/en/stable/distributions.html#foldeddistribution)

[TuanNguyen27]

TFP also has Geometric & Weibull distribution support, in case #710 is successful :)

[fehiepsi]

Yeah, if those distributions in numpyro.contrib.tfp.distributions work well, we will probably move them to the main module. There are some interesting distributions in tfp such as JointDistributionSequential.

[nkaimcaudle]

Is it possible to also add SinhArcsinh transformation ?

[fehiepsi]

@nkaimcaudle Sure, I just created a FR for it. :)

[dirmeier]

Hello all,

I'm working on two models with ZINB and Weibull likelihoods right now, if you are interested to have implementations for both of these, I could open up a PR.

Maybe it could also make sense to have the alternative parameterization of the NB (https://mc-stan.org/docs/2_26/functions-reference/nbalt.html), since that is the one mainly used for regression models.

Thank you!

Best,
Simon

[fehiepsi]

Hi @dirmeier, your proposals sound great to me. Regarding NB, you could name the other parameterization NegativeBinomial2. :)