Repository: pyro-ppl/numpyro

Issue:
Implement some popular distributions
This issue tracks some popular distributions that are not available in NumPyro yet.

- [x] [Geometric distribution](https://pytorch.org/docs/master/distributions.html#geometric)
- [x] [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)


Discussion:
TFP also has [Geometric ](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/Geometric) & [Weibull](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/Weibull) distribution support, in case #710 is successful :) 
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](https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/JointDistributionSequential).
Is it possible to also add [SinhArcsinh transformation](https://www.tensorflow.org/probability/api_docs/python/tfp/bijectors/SinhArcsinh) ?
@nkaimcaudle Sure, I just created a FR for it. :)
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
Hi @dirmeier, your proposals sound great to me. Regarding NB, you could name the other parameterization `NegativeBinomial2`. :)

Resolving patch:
diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
--- a/numpyro/distributions/__init__.py
+++ b/numpyro/distributions/__init__.py
@@ -64,6 +64,7 @@
     Delta,
     Distribution,
     ExpandedDistribution,
+    FoldedDistribution,
     ImproperUniform,
     Independent,
     MaskedDistribution,
@@ -109,6 +110,7 @@
     "Distribution",
     "Exponential",
     "ExpandedDistribution",
+    "FoldedDistribution",
     "Gamma",
     "GammaPoisson",
     "GaussianRandomWalk",
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -35,8 +35,9 @@
 
 from jax import lax, tree_util
 import jax.numpy as jnp
+from jax.scipy.special import logsumexp
 
-from numpyro.distributions.transforms import ComposeTransform, Transform
+from numpyro.distributions.transforms import AbsTransform, ComposeTransform, Transform
 from numpyro.distributions.util import (
     lazy_property,
     promote_shapes,
@@ -1051,6 +1052,38 @@ def tree_flatten(self):
         )
 
 
+class FoldedDistribution(TransformedDistribution):
+    """
+    Equivalent to ``TransformedDistribution(base_dist, AbsTransform())``,
+    but additionally supports :meth:`log_prob` .
+
+    :param Distribution base_dist: A univariate distribution to reflect.
+    """
+
+    support = constraints.positive
+
+    def __init__(self, base_dist, validate_args=None):
+        if base_dist.event_shape:
+            raise ValueError("Only univariate distributions can be folded.")
+        super().__init__(base_dist, AbsTransform(), validate_args=validate_args)
+
+    @validate_sample
+    def log_prob(self, value):
+        dim = max(len(self.batch_shape), jnp.ndim(value))
+        plus_minus = jnp.array([1.0, -1.0]).reshape((2,) + (1,) * dim)
+        return logsumexp(self.base_dist.log_prob(plus_minus * value), axis=0)
+
+    def tree_flatten(self):
+        base_flatten, base_aux = self.base_dist.tree_flatten()
+        return base_flatten, (type(self.base_dist), base_aux)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        base_cls, base_aux = aux_data
+        base_dist = base_cls.tree_unflatten(base_aux, params)
+        return cls(base_dist)
+
+
 class Delta(Distribution):
     arg_constraints = {
         "v": constraints.dependent(is_discrete=False),


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -109,6 +109,19 @@ def __init__(self, rate, *, validate_args=None):
         super().__init__(rate, is_sparse=True, validate_args=validate_args)
 
 
+class FoldedNormal(dist.FoldedDistribution):
+    arg_constraints = {"loc": constraints.real, "scale": constraints.positive}
+
+    def __init__(self, loc, scale, validate_args=None):
+        self.loc = loc
+        self.scale = scale
+        super().__init__(dist.Normal(loc, scale), validate_args=validate_args)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        return dist.FoldedDistribution.tree_unflatten(aux_data, params)
+
+
 _DIST_MAP = {
     dist.BernoulliProbs: lambda probs: osp.bernoulli(p=probs),
     dist.BernoulliLogits: lambda logits: osp.bernoulli(p=_to_probs_bernoulli(logits)),
@@ -189,6 +202,8 @@ def get_sp_dist(jax_dist):
     T(dist.Gumbel, 0.0, 1.0),
     T(dist.Gumbel, 0.5, 2.0),
     T(dist.Gumbel, jnp.array([0.0, 0.5]), jnp.array([1.0, 2.0])),
+    T(FoldedNormal, 2.0, 4.0),
+    T(FoldedNormal, jnp.array([2.0, 50.0]), jnp.array([4.0, 100.0])),
     T(dist.HalfCauchy, 1.0),
     T(dist.HalfCauchy, jnp.array([1.0, 2.0])),
     T(dist.HalfNormal, 1.0),
@@ -1102,6 +1117,8 @@ def fn(*args):
 def test_mean_var(jax_dist, sp_dist, params):
     if jax_dist is _ImproperWrapper:
         pytest.skip("Improper distribution does not has mean/var implemented")
+    if jax_dist is FoldedNormal:
+        pytest.skip("Folded distribution does not has mean/var implemented")
     if jax_dist in (
         _TruncatedNormal,
         dist.LeftTruncatedDistribution,