Repository: pyro-ppl/pyro

Issue:
FR an AsymmetricSoftLaplace distribution
It would be nice to have an `AsymmetricSoftLaplace` distribution, combining features of [AsymmetricLaplace](https://docs.pyro.ai/en/latest/distributions.html#asymmetriclaplace) and [SoftLaplace](https://docs.pyro.ai/en/latest/distributions.html#softlaplace):
- A (loc, scale) family amenable to `LocScaleReparam`
- Learnable skew
- Concave log density, leading to unimodal posteriors
- Infinitely differentiable, leading to smoother gradients in SVI and HMC and tighter Gaussian approximations

One approach is to use the density
```
p(x) = const / (exp(x/k) + exp(-k*x))
```
however according to [wolfram alpha](https://www.wolframalpha.com/input/?i=Integrate%5B1+%2F+%28e%5E%28x%2Fk%29+%2B+e%5E%28-k*x%29%29%2C+%7Bx%7D%5D), the normalizing `const` and `cdf` involve `torch.special.hyp2f1()` which is [not yet implemented](https://github.com/pytorch/pytorch/issues/50345).

Another approach might be to sum two Gamma random variables in opposite directions, with tunable smoothness (e.g. a shared `alpha` parameter). This has an easy `.rsample()` and I believe a tractable `.log_prob()`, but the density is not infinitely differentiable and requires a smoothness parameter.

A third possibility is to Gaussian-blur a standard `AsymmetricLaplace`.  This has tractable `.log_prob()` and `.rsample()` and a tunable smoothness parameter with a plausible default (i.e. half the variance is due to each of the asymmetric laplace and normal terms).


Discussion:

Resolving patch:
diff --git a/pyro/distributions/__init__.py b/pyro/distributions/__init__.py
--- a/pyro/distributions/__init__.py
+++ b/pyro/distributions/__init__.py
@@ -7,7 +7,10 @@
 # isort: split
 
 from pyro.distributions.affine_beta import AffineBeta
-from pyro.distributions.asymmetriclaplace import AsymmetricLaplace
+from pyro.distributions.asymmetriclaplace import (
+    AsymmetricLaplace,
+    SoftAsymmetricLaplace,
+)
 from pyro.distributions.avf_mvn import AVFMultivariateNormal
 from pyro.distributions.coalescent import (
     CoalescentRateLikelihood,
@@ -135,6 +138,7 @@
     "SineBivariateVonMises",
     "SineSkewed",
     "SoftLaplace",
+    "SoftAsymmetricLaplace",
     "SpanningTree",
     "Stable",
     "TorchDistribution",
diff --git a/pyro/distributions/asymmetriclaplace.py b/pyro/distributions/asymmetriclaplace.py
--- a/pyro/distributions/asymmetriclaplace.py
+++ b/pyro/distributions/asymmetriclaplace.py
@@ -1,6 +1,8 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+import math
+
 import torch
 from torch.distributions import constraints
 from torch.distributions.utils import broadcast_all, lazy_property
@@ -10,7 +12,8 @@
 
 class AsymmetricLaplace(TorchDistribution):
     """
-    Asymmetric version of Laplace distribution.
+    Asymmetric version of the :class:`~pyro.distributions.Laplace`
+    distribution.
 
     To the left of ``loc`` this acts like an
     ``-Exponential(1/(asymmetry*scale))``; to the right of ``loc`` this acts
@@ -27,7 +30,7 @@ class AsymmetricLaplace(TorchDistribution):
     support = constraints.real
     has_rsample = True
 
-    def __init__(self, loc, scale, asymmetry, validate_args=None):
+    def __init__(self, loc, scale, asymmetry, *, validate_args=None):
         self.loc, self.scale, self.asymmetry = broadcast_all(loc, scale, asymmetry)
         super().__init__(self.loc.shape, validate_args=validate_args)
 
@@ -74,3 +77,120 @@ def variance(self):
         p = left / total
         q = right / total
         return p * left ** 2 + q * right ** 2 + p * q * total ** 2
+
+
+class SoftAsymmetricLaplace(TorchDistribution):
+    """
+    Soft asymmetric version of the :class:`~pyro.distributions.Laplace`
+    distribution.
+
+    This has a smooth (infinitely differentiable) density with two asymmetric
+    asymptotically exponential tails, one on the left and one on the right. In
+    the limit of ``softness → 0``, this converges in distribution to the
+    :class:`AsymmetricLaplace` distribution.
+
+    This is equivalent to the sum of three random variables ``z - u + v`` where::
+
+        z ~ Normal(loc, scale * softness)
+        u ~ Exponential(1 / (scale * asymmetry))
+        v ~ Exponential(asymetry / scale)
+
+    This is also equivalent the sum of two random variables ``z + a`` where::
+
+        z ~ Normal(loc, scale * softness)
+        a ~ AsymmetricLaplace(0, scale, asymmetry)
+
+    :param loc: Location parameter, i.e. the mode.
+    :param scale: Scale parameter = geometric mean of left and right scales.
+    :param asymmetry: Square of ratio of left to right scales. Defaults to 1.
+    :param softness: Scale parameter of the Gaussian smoother. Defaults to 1.
+    """
+    arg_constraints = {"loc": constraints.real,
+                       "scale": constraints.positive,
+                       "asymmetry": constraints.positive,
+                       "softness": constraints.positive}
+    support = constraints.real
+    has_rsample = True
+
+    def __init__(self, loc, scale, asymmetry=1.0, softness=1.0, *, validate_args=None):
+        self.loc, self.scale, self.asymmetry, self.softness = broadcast_all(
+            loc, scale, asymmetry, softness,
+        )
+        super().__init__(self.loc.shape, validate_args=validate_args)
+
+    @lazy_property
+    def left_scale(self):
+        return self.scale * self.asymmetry
+
+    @lazy_property
+    def right_scale(self):
+        return self.scale / self.asymmetry
+
+    @lazy_property
+    def soft_scale(self):
+        return self.scale * self.softness
+
+    def expand(self, batch_shape, _instance=None):
+        new = self._get_checked_instance(AsymmetricLaplace, _instance)
+        batch_shape = torch.Size(batch_shape)
+        new.loc = self.loc.expand(batch_shape)
+        new.scale = self.scale.expand(batch_shape)
+        new.asymmetry = self.asymmetry.expand(batch_shape)
+        new.softness = self.softness.expand(batch_shape)
+        super(AsymmetricLaplace, new).__init__(batch_shape, validate_args=False)
+        new._validate_args = self._validate_args
+        return new
+
+    def log_prob(self, value):
+        if self._validate_args:
+            self._validate_sample(value)
+
+        # Standardize.
+        x = (value - self.loc) / self.scale
+        L = self.asymmetry
+        R = self.asymmetry.reciprocal()
+        S = self.softness
+        SS = S * S
+        S2 = S * math.sqrt(2)
+        Lx = L * x
+        Rx = R * x
+
+        # This is the sum of two integrals which are proportional to:
+        # left = Integrate[e^(-t/L - ((x+t)/S)^2/2)/sqrt(2 pi)/S, {t,0,Infinity}]
+        #      = 1/2 e^((2 L x + S^2)/(2 L^2)) erfc((L x + S^2)/(sqrt(2) L S))
+        # right = Integrate[e^(-t/R - ((x-t)/S)^2/2)/sqrt(2 pi)/S, {t,0,Infinity}]
+        #       = 1/2 e^((S^2 - 2 R x)/(2 R^2)) erfc((S^2 - R x)/(sqrt(2) R S))
+        return math.log(0.5) + torch.logaddexp(
+            (SS / 2 + Lx) / L ** 2 + _logerfc((SS + Lx) / (L * S2)),
+            (SS / 2 - Rx) / R ** 2 + _logerfc((SS - Rx) / (R * S2)),
+        ) - (L + R).log() - self.scale.log()
+
+    def rsample(self, sample_shape=torch.Size()):
+        shape = self._extended_shape(sample_shape)
+        z = self.loc.new_empty(shape).normal_()
+        u, v = self.loc.new_empty((2,) + shape).exponential_()
+        return (self.loc + self.soft_scale * z - self.left_scale * u
+                + self.right_scale * v)
+
+    @property
+    def mean(self):
+        total_scale = self.left_scale + self.right_scale
+        return self.loc + (self.right_scale ** 2 - self.left_scale ** 2) / total_scale
+
+    @property
+    def variance(self):
+        left = self.left_scale
+        right = self.right_scale
+        total = left + right
+        p = left / total
+        q = right / total
+        return (p * left ** 2 + q * right ** 2 + p * q * total ** 2
+                + self.soft_scale ** 2)
+
+
+def _logerfc(x):
+    try:
+        # Requires https://github.com/pytorch/pytorch/issues/31945
+        return torch.logerfc(x)
+    except AttributeError:
+        return x.double().erfc().log().to(dtype=x.dtype)


Test patch:
diff --git a/tests/distributions/conftest.py b/tests/distributions/conftest.py
--- a/tests/distributions/conftest.py
+++ b/tests/distributions/conftest.py
@@ -379,6 +379,13 @@ def __init__(self, von_loc, von_conc, skewness):
                 {'loc': [2.0, -50.0], 'scale': [2.0, 10.0],
                  'asymmetry': [0.5, 2.5], 'test_data': [[2.0, 10.0], [-1.0, -50.0]]},
             ]),
+    Fixture(pyro_dist=dist.SoftAsymmetricLaplace,
+            examples=[
+                {'loc': [1.0], 'scale': [1.0], 'asymmetry': [2.0],
+                 'test_data': [2.0]},
+                {'loc': [2.0, -50.0], 'scale': [2.0, 10.0], 'asymmetry': [0.5, 2.5],
+                 'softness': [0.7, 1.4], 'test_data': [[2.0, 10.0], [-1.0, -50.0]]},
+            ]),
 ]
 
 discrete_dists = [