From ecefe32045d3d004c0b804dfd4de93d437d0621c Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Fri, 12 Feb 2021 01:32:18 -0600
Subject: [PATCH 1/7] add truncated distribution and cdf/icdf method for some
 distributions

---
 numpyro/distributions/continuous.py | 162 +++++++++++++++++++++++++++-
 1 file changed, 161 insertions(+), 1 deletion(-)

diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index 3f6dbc5dc..8e5b3d3a5 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -31,7 +31,7 @@
 import jax.numpy as jnp
 import jax.random as random
 from jax.scipy.linalg import cho_solve, solve_triangular
-from jax.scipy.special import gammaln, log_ndtr, logsumexp, multigammaln, ndtr, ndtri
+from jax.scipy.special import expit, gammaln, logit, log_ndtr, logsumexp, multigammaln, ndtr, ndtri
 
 from numpyro.distributions import constraints
 from numpyro.distributions.distribution import Distribution, TransformedDistribution
@@ -109,6 +109,13 @@ def mean(self):
     def variance(self):
         return jnp.full(self.batch_shape, jnp.nan)
 
+    def cdf(self, value):
+        scaled = (value - self.loc) / self.scale
+        return jnp.arctan(scaled) / jnp.pi + 0.5
+
+    def icdf(self, q):
+        return self.loc + self.scale * jnp.tan(jnp.pi(q - 0.5))
+
 
 class Dirichlet(Distribution):
     arg_constraints = {'concentration': constraints.independent(constraints.positive, 1)}
@@ -412,6 +419,14 @@ def mean(self):
     def variance(self):
         return jnp.broadcast_to(2 * self.scale ** 2, self.batch_shape)
 
+    def cdf(self, value):
+        scaled = (value - self.loc) / self.scale
+        return 0.5 - 0.5 * jnp.sign(scaled) * jnp.expm1(-jnp.abs(scaled))
+
+    def icdf(self, q):
+        a = q - 0.5
+        return self.loc - self.scale * self.sign(a) * jnp.log1p(-2 * jnp.abs(a))
+
 
 class LKJ(TransformedDistribution):
     r"""
@@ -966,6 +981,10 @@ def log_prob(self, value):
         value_scaled = (value - self.loc) / self.scale
         return -0.5 * value_scaled ** 2 - normalize_term
 
+    def cdf(self, value):
+        scaled = (value - self.loc) / self.scale
+        return ndtr(scaled)
+
     def icdf(self, q):
         return self.loc + self.scale * ndtri(q)
 
@@ -1049,6 +1068,140 @@ def variance(self):
         var = jnp.where(self.df <= 1, jnp.nan, var)
         return jnp.broadcast_to(var, self.batch_shape)
 
+    def cdf(self, value):
+        pass  # TODO
+
+    def icdf(self, q):
+        pass  # TODO
+
+
+class OneSidedTruncatedDistribution(Distribution):
+    arg_constraints = {"low": constraints.real}
+    reparametrized_params = ["low"]
+
+    def __init__(self, base_dist, low=0., validate_args=None):
+        assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
+        assert base_dist.support is constraints.real, \
+            "The base distribution should be univariate and have real support."
+        batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low))
+        self.low = promote_shapes(low, batch_shape)
+        self._support = constraints.greater_than(low)
+        super().__init__(batch_shape, validate_args=validate_args)
+
+    @constraints.dependent_property(is_discrete=False, event_dim=0)
+    def support(self):
+        return self._support
+
+    @lazy_property
+    def _ccdf_low(self):
+        return self.base_dist.cdf(-self.low)
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        u = random.uniform(key, sample_shape + self.batch_shape)
+        # NB: we use a more numerical formula for a symmetric base distribution and non-negative low
+        #   icdf(cdf(low) + (1 - cdf(low)) * u) = -icdf[(1 - u) * (1 - cdf(low))]
+        #                                       = -icdf[(1 - u) * cdf(-low)]
+        return -self.base_dist.icdf(self._ccdf_low * (1 - u))
+
+    @validate_sample
+    def log_prob(self, value):
+        # NB: we use a more numerical formula for a symmetric base distribution and non-genative low
+        #   log(1 - cdf(low)) = logcdf(-low)
+        return self.base_dist.log_prob(value) - self._ccdf_low
+
+    def tree_flatten(self):
+        base_flatten, base_aux = self.base_dist.tree_flatten()
+        if isinstance(self._support.lower_bound, (int, float)):
+            return base_flatten, (type(self.base_dist), base_aux, self._support.lower_bound)
+        else:
+            return (base_flatten, self._support.lower_bound), (type(self.base_dist), base_aux)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        if len(aux_data) == 2:
+            base_flatten, low = params
+            base_cls, base_aux = aux_data
+        else:
+            base_flatten = params
+            base_cls, base_aux, low = aux_data
+        base_dist = base_cls.tree_unflatten(base_aux, base_flatten)
+        return cls(base_dist, low=low)
+
+
+class TwoSidedTruncatedDistribution(Distribution):
+    arg_constraints = {"low": constraints.dependent, "high": constraints.dependent}
+    reparametrized_params = ["low", "high"]
+
+    def __init__(self, base_dist, low=0., high=1., validate_args=None):
+        assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
+        assert base_dist.support is constraints.real, \
+            "The base distribution should be univariate and have real support."
+        batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low), jnp.shape(high))
+        self.low = promote_shapes(low, batch_shape)
+        self.high = promote_shapes(high, batch_shape)
+        self._support = constraints.interval(low, high)
+        super().__init__(batch_shape, validate_args=validate_args)
+
+    @constraints.dependent_property(is_discrete=False, event_dim=0)
+    def support(self):
+        return self._support
+
+    @lazy_property
+    def _ccdf_low(self):
+        return self.base_dist.cdf(-self.low)
+
+    @lazy_property
+    def _ccdf_high(self):
+        return self.base_dist.cdf(-self.high)
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        u = random.uniform(key, sample_shape + self.batch_shape)
+        return -self.base_dist.icdf(self._ccdf_low * (1 - u) + self._ccdf_high * u)
+
+    @validate_sample
+    def log_prob(self, value):
+        return self.base_dist.log_prob(value) - jnp.log(self._ccdf_low - self._ccdf_high)
+
+    def tree_flatten(self):
+        base_flatten, base_aux = self.base_dist.tree_flatten()
+        if isinstance(self._support.lower_bound, (int, float)) and \
+                isinstance(self._support.upper_bound, (int, float)):
+            return base_flatten, (type(self.base_dist), base_aux,
+                                  self._support.lower_bound, self._support.upper_bound)
+        else:
+            return (base_flatten, self._support.lower_bound, self.support.upper_bound), \
+                (type(self.base_dist), base_aux)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        if len(aux_data) == 2:
+            base_flatten, low, high = params
+            base_cls, base_aux = aux_data
+        else:
+            base_flatten = params
+            base_cls, base_aux, low, high = aux_data
+        base_dist = base_cls.tree_unflatten(base_aux, base_flatten)
+        return cls(base_dist, low=low, high=high)
+
+
+def TruncatedDistribution(base_dist, low=0., high=None, validate_args=None):
+    """
+    If `high` is None, this is a one-sided truncated distribution. Otherwise,
+    this is a two-sided truncated distribution.
+
+    Currently, this class only supports truncating Normal, Cauchy, or StudentT distributions.
+
+    .. note:: The implementation is customized to be more stable when `low` is non-negative.
+        For small negative parameters (says `low < -5`), this distribution might return
+        non-finite values.
+    """
+    if high is None:
+        return OneSidedTruncatedDistribution(base_dist, low=low, validate_args=validate_args)
+    else:
+        return TwoSidedTruncatedDistribution(base_dist, low=low, high=high, validate_args=validate_args)
+
 
 class _BaseTruncatedCauchy(Distribution):
     # NB: this is a truncated cauchy with low=0, scale=1
@@ -1282,6 +1435,13 @@ def variance(self):
         var = (self.scale ** 2) * (jnp.pi ** 2) / 3
         return jnp.broadcast_to(var, self.batch_shape)
 
+    def cdf(self, value):
+        scaled = (value - self.loc) / self.scale
+        return expit(scaled)
+
+    def icdf(self, q):
+        return self.loc + self.scale * logit(q)
+
 
 class TruncatedPolyaGamma(Distribution):
     truncation_point = 2.5

From cf5edbd252f1715a5c3fdb37fde726aeb9e1703b Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Fri, 12 Feb 2021 15:45:47 -0600
Subject: [PATCH 2/7] add student-t cdf and numerical issues

---
 numpyro/distributions/continuous.py | 120 ++++++++++++++++++++++------
 1 file changed, 97 insertions(+), 23 deletions(-)

diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index 8e5b3d3a5..1032bfa88 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -31,7 +31,7 @@
 import jax.numpy as jnp
 import jax.random as random
 from jax.scipy.linalg import cho_solve, solve_triangular
-from jax.scipy.special import expit, gammaln, logit, log_ndtr, logsumexp, multigammaln, ndtr, ndtri
+from jax.scipy.special import betainc, expit, gammaln, logit, log_ndtr, logsumexp, multigammaln, ndtr, ndtri
 
 from numpyro.distributions import constraints
 from numpyro.distributions.distribution import Distribution, TransformedDistribution
@@ -81,6 +81,9 @@ def variance(self):
         total = self.concentration1 + self.concentration0
         return self.concentration1 * self.concentration0 / (total ** 2 * (total + 1))
 
+    def cdf(self, value):
+        return betainc(self.concentration1, self.concentration0, value)
+
 
 class Cauchy(Distribution):
     arg_constraints = {'loc': constraints.real, 'scale': constraints.positive}
@@ -1069,13 +1072,22 @@ def variance(self):
         return jnp.broadcast_to(var, self.batch_shape)
 
     def cdf(self, value):
-        pass  # TODO
+        # Ref: https://en.wikipedia.org/wiki/Student's_t-distribution#Related_distributions
+        # X^2 ~ F(1, df) -> df / (df + X^2) ~ Beta(df/2, 0.5)
+        scaled = (value - self.loc) / self.scale
+        scaled_squared = scaled * scaled
+        beta_value = self.df / (self.df + scaled_squared)
+        # when scaled < 0, returns 0.5 * Beta(df/2, 0.5).cdf(beta_value)
+        # when scaled > 0, returns 1 - 0.5 * Beta(df/2, 0.5).cdf(beta_value)
+        return 0.5 * (1 + jnp.sign(scaled) * (1 - betainc(0.5 * self.df, 0.5, beta_value)))
 
     def icdf(self, q):
-        pass  # TODO
+        # upstream issue: https://github.com/google/jax/issues/2399
+        raise NotImplementedError("Not implemented until scipy.special.betaincinv is"
+                                  " avaiable in JAX.")
 
 
-class OneSidedTruncatedDistribution(Distribution):
+class LeftTruncatedDistribution(Distribution):
     arg_constraints = {"low": constraints.real}
     reparametrized_params = ["low"]
 
@@ -1092,23 +1104,31 @@ def __init__(self, base_dist, low=0., validate_args=None):
     def support(self):
         return self._support
 
-    @lazy_property
-    def _ccdf_low(self):
-        return self.base_dist.cdf(-self.low)
-
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
-        # NB: we use a more numerical formula for a symmetric base distribution and non-negative low
+        # NB: we use a more numerical formula for a symmetric base distribution
+        # if low < loc
+        #   icdf(cdf(low) + (1 - cdf(low)) * u) =  icdf[(1 - u) * cdf(low) + u]
+        # if low > loc
         #   icdf(cdf(low) + (1 - cdf(low)) * u) = -icdf[(1 - u) * (1 - cdf(low))]
-        #                                       = -icdf[(1 - u) * cdf(-low)]
-        return -self.base_dist.icdf(self._ccdf_low * (1 - u))
+        #                                       = -icdf[(1 - u) * cdf(2*loc-low)]
+        diff = self.loc - self.low
+        cdf = self.base_dist.cdf(self.loc - jnp.abs(diff))
+        sign = jnp.where(diff >= 0, 1., -1.)
+        return sign * self.base_dist.icdf((1 - u) * cdf + 0.5 * (1 + sign) * u)
 
     @validate_sample
     def log_prob(self, value):
-        # NB: we use a more numerical formula for a symmetric base distribution and non-genative low
-        #   log(1 - cdf(low)) = logcdf(-low)
-        return self.base_dist.log_prob(value) - self._ccdf_low
+        # NB: we use a more numerical formula for a symmetric base distribution
+        # if low < loc
+        #   1 - cdf(low) = as-is
+        # if low > loc
+        #   1 - cdf(low) = cdf(2 * loc - low)
+        diff = self.loc - self.low
+        cdf = self.base_dist.cdf(self.loc - jnp.abs(diff))
+        sign = jnp.where(diff >= 0, 1., -1.)
+        return self.base_dist.log_prob(value) - jnp.log(0.5 + sign * (0.5 - cdf))
 
     def tree_flatten(self):
         base_flatten, base_aux = self.base_dist.tree_flatten()
@@ -1129,6 +1149,41 @@ def tree_unflatten(cls, aux_data, params):
         return cls(base_dist, low=low)
 
 
+class RightTruncatedDistribution(TransformedDistribution):
+    arg_constraints = {"high": constraints.real}
+    reparametrized_params = ["high"]
+
+    def __init__(self, base_dist, high=0., validate_args=None):
+        assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
+        loc2 = 2 * base_dist.loc
+        low = loc2 - high
+        left_truncated_dist = LeftTruncatedDistribution(base_dist, low=low, validate_args=validate_args)
+        super().__init__(left_truncated_dist, AffineTransform(loc2, -1), validate_args=validate_args)
+        self._support = constraints.less_than(high)
+
+    @constraints.dependent_property(is_discrete=False, event_dim=0)
+    def support(self):
+        return self._support
+
+    def tree_flatten(self):
+        base_flatten, base_aux = self.base_dist.base_dist.tree_flatten()
+        if isinstance(self._support.upper_bound, (int, float)):
+            return base_flatten, (type(self.base_dist), base_aux, self._support.upper_bound)
+        else:
+            return (base_flatten, self._support.upper_bound), (type(self.base_dist), base_aux)
+
+    @classmethod
+    def tree_unflatten(cls, aux_data, params):
+        if len(aux_data) == 2:
+            base_flatten, high = params
+            base_cls, base_aux = aux_data
+        else:
+            base_flatten = params
+            base_cls, base_aux, high = aux_data
+        base_dist = base_cls.tree_unflatten(base_aux, base_flatten)
+        return cls(base_dist, high=high)
+
+
 class TwoSidedTruncatedDistribution(Distribution):
     arg_constraints = {"low": constraints.dependent, "high": constraints.dependent}
     reparametrized_params = ["low", "high"]
@@ -1158,6 +1213,19 @@ def _ccdf_high(self):
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
+
+        # NB: we use a more numerical formula for a symmetric base distribution
+        #   A = icdf(cdf(low) + (cdf(high) - cdf(low)) * u) =  icdf[(1 - u) * cdf(low) + u * cdf(high)]
+        # will suffer by precision issues when low is large or when high is small;
+        # so we will split the implementation into 3 cases
+        #   low > loc, high < loc, and (loc <= loc or high >= loc)
+        # If low > loc:
+        #   A = ...
+        # Elif high < loc:
+        #   A = ...
+        # Otherwise,
+        #   ...
+
         return -self.base_dist.icdf(self._ccdf_low * (1 - u) + self._ccdf_high * u)
 
     @validate_sample
@@ -1188,17 +1256,23 @@ def tree_unflatten(cls, aux_data, params):
 
 def TruncatedDistribution(base_dist, low=0., high=None, validate_args=None):
     """
-    If `high` is None, this is a one-sided truncated distribution. Otherwise,
-    this is a two-sided truncated distribution.
-
-    Currently, this class only supports truncating Normal, Cauchy, or StudentT distributions.
-
-    .. note:: The implementation is customized to be more stable when `low` is non-negative.
-        For small negative parameters (says `low < -5`), this distribution might return
-        non-finite values.
+    A function to generate a truncated distribution.
+
+    :param base_dist: The base distribution to be truncated. This should be a univariate
+        distribution. Currently, only the following distributions are supported:
+        Cauchy, Laplace, Logistic, Normal, and StudentT.
+    :param low: the value which is used to truncate the base distribution from below.
+        Setting this parameter to None to not truncate from below.
+    :param high: the value which is used to truncate the base distribution from above.
+        Setting this parameter to None to not truncate from above.
     """
     if high is None:
-        return OneSidedTruncatedDistribution(base_dist, low=low, validate_args=validate_args)
+        if low is None:
+            return base_dist
+        else:
+            return LeftTruncatedDistribution(base_dist, low=low, validate_args=validate_args)
+    elif low is None:
+        return RightTruncatedDistribution(base_dist, high=high, validate_args=validate_args)
     else:
         return TwoSidedTruncatedDistribution(base_dist, low=low, high=high, validate_args=validate_args)
 

From 9b2eaf3f07f60afe7d8fe699f9c5a604fe98fb39 Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Fri, 12 Feb 2021 20:54:18 -0600
Subject: [PATCH 3/7] fix the implementation

---
 docs/source/distributions.rst         |  8 ++++
 numpyro/distributions/__init__.py     |  2 +
 numpyro/distributions/continuous.py   | 57 ++++++++++++++-------------
 numpyro/distributions/distribution.py | 18 +++++++++
 4 files changed, 58 insertions(+), 27 deletions(-)

diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
index 5ab0222ef..47bca8ac9 100644
--- a/docs/source/distributions.rst
+++ b/docs/source/distributions.rst
@@ -245,6 +245,14 @@ TruncatedCauchy
     :show-inheritance:
     :member-order: bysource
 
+TruncatedDistribution
+---------------------
+.. autoclass:: numpyro.distributions.continuous.TruncatedDistribution
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 TruncatedNormal
 ---------------
 .. autoclass:: numpyro.distributions.continuous.TruncatedNormal
diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
index cc5fbec9a..9f2cb662c 100644
--- a/numpyro/distributions/__init__.py
+++ b/numpyro/distributions/__init__.py
@@ -26,6 +26,7 @@
     Pareto,
     StudentT,
     TruncatedCauchy,
+    TruncatedDistribution,
     TruncatedNormal,
     TruncatedPolyaGamma,
     Uniform
@@ -121,6 +122,7 @@
     'StudentT',
     'TransformedDistribution',
     'TruncatedCauchy',
+    'TruncatedDistribution',
     'TruncatedNormal',
     'TruncatedPolyaGamma',
     'Uniform',
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index 1032bfa88..28be58c23 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -1111,12 +1111,13 @@ def sample(self, key, sample_shape=()):
         # if low < loc
         #   icdf(cdf(low) + (1 - cdf(low)) * u) =  icdf[(1 - u) * cdf(low) + u]
         # if low > loc
-        #   icdf(cdf(low) + (1 - cdf(low)) * u) = -icdf[(1 - u) * (1 - cdf(low))]
-        #                                       = -icdf[(1 - u) * cdf(2*loc-low)]
+        #   icdf(cdf(low) + (1 - cdf(low)) * u) = loc - icdf[(1 - u) * (1 - cdf(low))]
+        #                                       = loc - icdf[(1 - u) * cdf(2*loc-low)]
         diff = self.loc - self.low
-        cdf = self.base_dist.cdf(self.loc - jnp.abs(diff))
         sign = jnp.where(diff >= 0, 1., -1.)
-        return sign * self.base_dist.icdf((1 - u) * cdf + 0.5 * (1 + sign) * u)
+        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
+        high_cdf = 0.5 * (1 + sign)
+        return 0.5 * (1 - sign) * self.loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
 
     @validate_sample
     def log_prob(self, value):
@@ -1126,9 +1127,10 @@ def log_prob(self, value):
         # if low > loc
         #   1 - cdf(low) = cdf(2 * loc - low)
         diff = self.loc - self.low
-        cdf = self.base_dist.cdf(self.loc - jnp.abs(diff))
         sign = jnp.where(diff >= 0, 1., -1.)
-        return self.base_dist.log_prob(value) - jnp.log(0.5 + sign * (0.5 - cdf))
+        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
+        high_cdf = 0.5 * (1 + sign)
+        return self.base_dist.log_prob(value) - jnp.log(sign * (high_cdf - low_cdf))
 
     def tree_flatten(self):
         base_flatten, base_aux = self.base_dist.tree_flatten()
@@ -1202,35 +1204,36 @@ def __init__(self, base_dist, low=0., high=1., validate_args=None):
     def support(self):
         return self._support
 
-    @lazy_property
-    def _ccdf_low(self):
-        return self.base_dist.cdf(-self.low)
-
-    @lazy_property
-    def _ccdf_high(self):
-        return self.base_dist.cdf(-self.high)
-
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
 
         # NB: we use a more numerical formula for a symmetric base distribution
-        #   A = icdf(cdf(low) + (cdf(high) - cdf(low)) * u) =  icdf[(1 - u) * cdf(low) + u * cdf(high)]
-        # will suffer by precision issues when low is large or when high is small;
-        # so we will split the implementation into 3 cases
-        #   low > loc, high < loc, and (loc <= loc or high >= loc)
-        # If low > loc:
-        #   A = ...
-        # Elif high < loc:
-        #   A = ...
-        # Otherwise,
-        #   ...
-
-        return -self.base_dist.icdf(self._ccdf_low * (1 - u) + self._ccdf_high * u)
+        #   A = icdf(cdf(low) + (cdf(high) - cdf(low)) * u) = icdf[(1 - u) * cdf(low) + u * cdf(high)]
+        # will suffer by precision issues when low is large;
+        # If low < loc:
+        #   A = icdf[(1 - u) * cdf(low) + u * cdf(high)]
+        # Else
+        #   A = loc - icdf[(1 - u) * cdf(2*loc-low)) + u * cdf(2*loc - high)]
+        # TODO: cache sign, low_cdf, high_cdf
+        diff = self.loc - self.low
+        sign = jnp.where(diff >= 0, 1., -1.)
+        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
+        high_cdf = self.base_dist.cdf(self.loc - sign * (self.loc - self.high))
+        return 0.5 * (1 - sign) * self.loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
 
     @validate_sample
     def log_prob(self, value):
-        return self.base_dist.log_prob(value) - jnp.log(self._ccdf_low - self._ccdf_high)
+        # NB: we use a more numerical formula for a symmetric base distribution
+        # if low < loc
+        #   cdf(high) - cdf(low) = as-is
+        # if low > loc
+        #   cdf(high) - cdf(low) = cdf(2 * loc - low) - cdf(2 * loc - high)
+        diff = self.loc - self.low
+        sign = jnp.where(diff >= 0, 1., -1.)
+        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
+        high_cdf = self.base_dist.cdf(self.loc - sign * (self.loc - self.high))
+        return self.base_dist.log_prob(value) - jnp.log(sign * (high_cdf - low_cdf))
 
     def tree_flatten(self):
         base_flatten, base_aux = self.base_dist.tree_flatten()
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
index 599424af0..898f093c0 100644
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -391,6 +391,24 @@ def infer_shapes(cls, *args, **kwargs):
         event_shape = ()
         return batch_shape, event_shape
 
+    def cdf(self, value):
+        """
+        The cummulative distribution function of this distribution.
+
+        :param value: samples from this distribution.
+        :return: output of the cummulative distribution function evaluated at `value`.
+        """
+        raise NotImplementedError
+
+    def icdf(self, q):
+        """
+        The inverse cumulative distribution function of this distribution.
+
+        :param q: quantile values, should belong to [0, 1].
+        :return: the samples whose cdf values equals to `q`.
+        """
+        raise NotImplementedError
+
 
 class ExpandedDistribution(Distribution):
     arg_constraints = {}

From 4584648e8362730bd2964351ac6901afe5607367 Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Fri, 12 Feb 2021 23:11:17 -0600
Subject: [PATCH 4/7] gof detects a buggit add -u .!

---
 numpyro/distributions/continuous.py | 53 ++++++++++++++++-------------
 test/test_distributions.py          | 50 ++++++++++++++++++++++++++-
 2 files changed, 79 insertions(+), 24 deletions(-)

diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index 2098145a3..c8a343bba 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -117,7 +117,7 @@ def cdf(self, value):
         return jnp.arctan(scaled) / jnp.pi + 0.5
 
     def icdf(self, q):
-        return self.loc + self.scale * jnp.tan(jnp.pi(q - 0.5))
+        return self.loc + self.scale * jnp.tan(jnp.pi * (q - 0.5))
 
 
 class Dirichlet(Distribution):
@@ -428,7 +428,7 @@ def cdf(self, value):
 
     def icdf(self, q):
         a = q - 0.5
-        return self.loc - self.scale * self.sign(a) * jnp.log1p(-2 * jnp.abs(a))
+        return self.loc - self.scale * jnp.sign(a) * jnp.log1p(-2 * jnp.abs(a))
 
 
 class LKJ(TransformedDistribution):
@@ -1082,9 +1082,9 @@ def cdf(self, value):
         return 0.5 * (1 + jnp.sign(scaled) * (1 - betainc(0.5 * self.df, 0.5, beta_value)))
 
     def icdf(self, q):
+        # scipy.special.betaincinv is not avaiable yet in JAX
         # upstream issue: https://github.com/google/jax/issues/2399
-        raise NotImplementedError("Not implemented until scipy.special.betaincinv is"
-                                  " avaiable in JAX.")
+        raise NotImplementedError
 
 
 class LeftTruncatedDistribution(Distribution):
@@ -1096,7 +1096,8 @@ def __init__(self, base_dist, low=0., validate_args=None):
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low))
-        self.low = promote_shapes(low, batch_shape)
+        self.base_dist = base_dist
+        self.low, = promote_shapes(low, shape=batch_shape)
         self._support = constraints.greater_than(low)
         super().__init__(batch_shape, validate_args=validate_args)
 
@@ -1111,13 +1112,14 @@ def sample(self, key, sample_shape=()):
         # if low < loc
         #   icdf(cdf(low) + (1 - cdf(low)) * u) =  icdf[(1 - u) * cdf(low) + u]
         # if low > loc
-        #   icdf(cdf(low) + (1 - cdf(low)) * u) = loc - icdf[(1 - u) * (1 - cdf(low))]
-        #                                       = loc - icdf[(1 - u) * cdf(2*loc-low)]
-        diff = self.loc - self.low
+        #   icdf(cdf(low) + (1 - cdf(low)) * u) = 2 * loc - icdf[(1 - u) * (1 - cdf(low))]
+        #                                       = 2 * loc - icdf[(1 - u) * cdf(2*loc-low)]
+        loc = self.base_dist.loc
+        diff = loc - self.low
         sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
+        low_cdf = self.base_dist.cdf(loc - sign * diff)
         high_cdf = 0.5 * (1 + sign)
-        return 0.5 * (1 - sign) * self.loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
+        return (1 - sign) * loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
 
     @validate_sample
     def log_prob(self, value):
@@ -1126,9 +1128,10 @@ def log_prob(self, value):
         #   1 - cdf(low) = as-is
         # if low > loc
         #   1 - cdf(low) = cdf(2 * loc - low)
-        diff = self.loc - self.low
+        loc = self.base_dist.loc
+        diff = loc - self.low
         sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
+        low_cdf = self.base_dist.cdf(loc - sign * diff)
         high_cdf = 0.5 * (1 + sign)
         return self.base_dist.log_prob(value) - jnp.log(sign * (high_cdf - low_cdf))
 
@@ -1161,6 +1164,7 @@ def __init__(self, base_dist, high=0., validate_args=None):
         low = loc2 - high
         left_truncated_dist = LeftTruncatedDistribution(base_dist, low=low, validate_args=validate_args)
         super().__init__(left_truncated_dist, AffineTransform(loc2, -1), validate_args=validate_args)
+        self.high, = promote_shapes(high, shape=self.batch_shape)
         self._support = constraints.less_than(high)
 
     @constraints.dependent_property(is_discrete=False, event_dim=0)
@@ -1195,8 +1199,9 @@ def __init__(self, base_dist, low=0., high=1., validate_args=None):
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low), jnp.shape(high))
-        self.low = promote_shapes(low, batch_shape)
-        self.high = promote_shapes(high, batch_shape)
+        self.base_dist = base_dist
+        self.low, = promote_shapes(low, shape=batch_shape)
+        self.high, = promote_shapes(high, shape=batch_shape)
         self._support = constraints.interval(low, high)
         super().__init__(batch_shape, validate_args=validate_args)
 
@@ -1214,13 +1219,14 @@ def sample(self, key, sample_shape=()):
         # If low < loc:
         #   A = icdf[(1 - u) * cdf(low) + u * cdf(high)]
         # Else
-        #   A = loc - icdf[(1 - u) * cdf(2*loc-low)) + u * cdf(2*loc - high)]
+        #   A = 2 * loc - icdf[(1 - u) * cdf(2*loc-low)) + u * cdf(2*loc - high)]
         # TODO: cache sign, low_cdf, high_cdf
-        diff = self.loc - self.low
+        loc = self.base_dist.loc
+        diff = loc - self.low
         sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
-        high_cdf = self.base_dist.cdf(self.loc - sign * (self.loc - self.high))
-        return 0.5 * (1 - sign) * self.loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
+        low_cdf = self.base_dist.cdf(loc - sign * diff)
+        high_cdf = self.base_dist.cdf(loc - sign * (loc - self.high))
+        return (1 - sign) * loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
 
     @validate_sample
     def log_prob(self, value):
@@ -1229,10 +1235,11 @@ def log_prob(self, value):
         #   cdf(high) - cdf(low) = as-is
         # if low > loc
         #   cdf(high) - cdf(low) = cdf(2 * loc - low) - cdf(2 * loc - high)
-        diff = self.loc - self.low
+        loc = self.base_dist.loc
+        diff = loc - self.low
         sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(self.loc - sign * diff)
-        high_cdf = self.base_dist.cdf(self.loc - sign * (self.loc - self.high))
+        low_cdf = self.base_dist.cdf(loc - sign * diff)
+        high_cdf = self.base_dist.cdf(loc - sign * (loc - self.high))
         return self.base_dist.log_prob(value) - jnp.log(sign * (high_cdf - low_cdf))
 
     def tree_flatten(self):
@@ -1257,7 +1264,7 @@ def tree_unflatten(cls, aux_data, params):
         return cls(base_dist, low=low, high=high)
 
 
-def TruncatedDistribution(base_dist, low=0., high=None, validate_args=None):
+def TruncatedDistribution(base_dist, low=None, high=None, validate_args=None):
     """
     A function to generate a truncated distribution.
 
diff --git a/test/test_distributions.py b/test/test_distributions.py
index d8b746c38..f9f982794 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -60,6 +60,26 @@ def _lowrank_mvn_to_scipy(loc, cov_fac, cov_diag):
     return osp.multivariate_normal(mean=mean, cov=cov)
 
 
+def _truncnorm_to_scipy(loc, scale, low, high):
+    if low is None:
+        a = -np.inf
+    else:
+        a = (low - loc) / scale
+    if high is None:
+        b = np.inf
+    else:
+        b = (high - loc) / scale
+    return osp.truncnorm(a, b, loc=loc, scale=scale)
+
+
+def _TruncatedNormal(loc, scale, low, high):
+    return dist.TruncatedDistribution(dist.Normal(loc, scale), low, high)
+
+
+_TruncatedNormal.reparametrized_params = []
+_TruncatedNormal.infer_shapes = lambda *args: (lax.broadcast_shapes(*args), ())
+
+
 class _ImproperWrapper(dist.ImproperUniform):
     def sample(self, key, sample_shape=()):
         transform = biject_to(self.support)
@@ -103,7 +123,8 @@ def sample(self, key, sample_shape=()):
     dist.Uniform: lambda a, b: osp.uniform(a, b - a),
     dist.Logistic: lambda loc, scale: osp.logistic(loc=loc, scale=scale),
     dist.VonMises: lambda loc, conc: osp.vonmises(loc=np.array(loc, dtype=np.float64),
-                                                  kappa=np.array(conc, dtype=np.float64))
+                                                  kappa=np.array(conc, dtype=np.float64)),
+    _TruncatedNormal: _truncnorm_to_scipy,
 }
 
 CONTINUOUS = [
@@ -178,6 +199,11 @@ def sample(self, key, sample_shape=()):
     T(dist.TruncatedNormal, -1., 0., 1.),
     T(dist.TruncatedNormal, 1., -1., jnp.array([1., 2.])),
     T(dist.TruncatedNormal, jnp.array([-2., 2.]), jnp.array([0., 1.]), jnp.array([[1.], [2.]])),
+    T(_TruncatedNormal, -1., 2., 1., 5.),
+    T(_TruncatedNormal, jnp.array([-1., 4.]), 2., None, 5.),
+    T(_TruncatedNormal, -1., jnp.array([2., 3.]), 1., None),
+    T(_TruncatedNormal, -1., 2., jnp.array([-6., 4.]), jnp.array([-4., 6.])),
+    T(_TruncatedNormal, jnp.array([0., 1.]), jnp.array([[1.], [2.]]), None, jnp.array([-2., 2.])),
     T(dist.Uniform, 0., 2.),
     T(dist.Uniform, 1., jnp.array([2., 3.])),
     T(dist.Uniform, jnp.array([0., 0.]), jnp.array([[2.], [3.]])),
@@ -519,6 +545,28 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
     assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5)
 
 
+@pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS)
+def test_cdf_and_icdf(jax_dist, sp_dist, params):
+    if not sp_dist:
+        pytest.skip('no corresponding scipy distn.')
+    jax_dist = jax_dist(*params)
+    if jax_dist.event_dim > 0:
+        pytest.skip('skip testing cdf/icdf methods of multivariate distributions')
+    sp_dist = sp_dist(*params)
+    rng_key = random.PRNGKey(0)
+    samples = jax_dist.sample(key=rng_key)
+    quantiles = random.uniform(random.PRNGKey(1), jax_dist.shape())
+    try:
+        actual_cdf = jax_dist.cdf(samples)
+        expected_cdf = sp_dist.cdf(samples)
+        assert_allclose(actual_cdf, expected_cdf, rtol=1e-5)
+        actual_icdf = jax_dist.icdf(quantiles)
+        expected_icdf = sp_dist.ppf(quantiles)
+        assert_allclose(actual_icdf, expected_icdf, rtol=1e-5)
+    except NotImplementedError:
+        pass
+
+
 @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS)
 def test_gof(jax_dist, sp_dist, params):
     if "Improper" in jax_dist.__name__:

From 92ab042418239b07afd66f1555dd77fdbcb6b070 Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Sat, 13 Feb 2021 00:59:44 -0600
Subject: [PATCH 5/7] fix failing tests

---
 numpyro/distributions/continuous.py | 115 ++++++++++++++++------------
 test/test_distributions.py          |  11 +++
 2 files changed, 78 insertions(+), 48 deletions(-)

diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index c8a343bba..26f74bb4f 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -26,7 +26,7 @@
 # POSSIBILITY OF SUCH DAMAGE.
 
 
-from jax import lax, ops
+from jax import lax, ops, tree_map
 import jax.nn as nn
 import jax.numpy as jnp
 import jax.random as random
@@ -1096,7 +1096,7 @@ def __init__(self, base_dist, low=0., validate_args=None):
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low))
-        self.base_dist = base_dist
+        self.base_dist = tree_map(lambda p: promote_shapes(p, shape=batch_shape)[0], base_dist)
         self.low, = promote_shapes(low, shape=batch_shape)
         self._support = constraints.greater_than(low)
         super().__init__(batch_shape, validate_args=validate_args)
@@ -1105,42 +1105,40 @@ def __init__(self, base_dist, low=0., validate_args=None):
     def support(self):
         return self._support
 
+    @lazy_property
+    def _tail_prob_at_low(self):
+        # if low < loc, returns cdf(low); otherwise returns 1 - cdf(low)
+        loc = self.base_dist.loc
+        sign = jnp.where(loc >= self.low, 1., -1.)
+        return self.base_dist.cdf(loc - sign * (loc - self.low))
+
+    @lazy_property
+    def _tail_prob_at_high(self):
+        # if low < loc, returns cdf(high) = 1; otherwise returns 1 - cdf(high) = 0
+        loc = self.base_dist.loc
+        sign = jnp.where(loc >= self.low, 1., -1.)
+        return 0.5 * (1 + sign)
+
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
-        # NB: we use a more numerical formula for a symmetric base distribution
-        # if low < loc
-        #   icdf(cdf(low) + (1 - cdf(low)) * u) =  icdf[(1 - u) * cdf(low) + u]
-        # if low > loc
-        #   icdf(cdf(low) + (1 - cdf(low)) * u) = 2 * loc - icdf[(1 - u) * (1 - cdf(low))]
-        #                                       = 2 * loc - icdf[(1 - u) * cdf(2*loc-low)]
         loc = self.base_dist.loc
-        diff = loc - self.low
-        sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(loc - sign * diff)
-        high_cdf = 0.5 * (1 + sign)
-        return (1 - sign) * loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
+        sign = jnp.where(loc >= self.low, 1., -1.)
+        return (1 - sign) * loc + sign * self.base_dist.icdf(
+            (1 - u) * self._tail_prob_at_low + u * self._tail_prob_at_high)
 
     @validate_sample
     def log_prob(self, value):
-        # NB: we use a more numerical formula for a symmetric base distribution
-        # if low < loc
-        #   1 - cdf(low) = as-is
-        # if low > loc
-        #   1 - cdf(low) = cdf(2 * loc - low)
-        loc = self.base_dist.loc
-        diff = loc - self.low
-        sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(loc - sign * diff)
-        high_cdf = 0.5 * (1 + sign)
-        return self.base_dist.log_prob(value) - jnp.log(sign * (high_cdf - low_cdf))
+        sign = jnp.where(self.base_dist.loc >= self.low, 1., -1.)
+        return self.base_dist.log_prob(value) - \
+            jnp.log(sign * (self._tail_prob_at_high - self._tail_prob_at_low))
 
     def tree_flatten(self):
         base_flatten, base_aux = self.base_dist.tree_flatten()
         if isinstance(self._support.lower_bound, (int, float)):
             return base_flatten, (type(self.base_dist), base_aux, self._support.lower_bound)
         else:
-            return (base_flatten, self._support.lower_bound), (type(self.base_dist), base_aux)
+            return (base_flatten, self.low), (type(self.base_dist), base_aux)
 
     @classmethod
     def tree_unflatten(cls, aux_data, params):
@@ -1154,29 +1152,43 @@ def tree_unflatten(cls, aux_data, params):
         return cls(base_dist, low=low)
 
 
-class RightTruncatedDistribution(TransformedDistribution):
+class RightTruncatedDistribution(Distribution):
     arg_constraints = {"high": constraints.real}
     reparametrized_params = ["high"]
 
     def __init__(self, base_dist, high=0., validate_args=None):
         assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
-        loc2 = 2 * base_dist.loc
-        low = loc2 - high
-        left_truncated_dist = LeftTruncatedDistribution(base_dist, low=low, validate_args=validate_args)
-        super().__init__(left_truncated_dist, AffineTransform(loc2, -1), validate_args=validate_args)
-        self.high, = promote_shapes(high, shape=self.batch_shape)
+        assert base_dist.support is constraints.real, \
+            "The base distribution should be univariate and have real support."
+        batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(high))
+        self.base_dist = tree_map(lambda p: promote_shapes(p, shape=batch_shape)[0], base_dist)
+        self.high, = promote_shapes(high, shape=batch_shape)
         self._support = constraints.less_than(high)
+        super().__init__(batch_shape, validate_args=validate_args)
 
     @constraints.dependent_property(is_discrete=False, event_dim=0)
     def support(self):
         return self._support
 
+    @lazy_property
+    def _cdf_at_high(self):
+        return self.base_dist.cdf(self.high)
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        u = random.uniform(key, sample_shape + self.batch_shape)
+        return self.base_dist.icdf(u * self._cdf_at_high)
+
+    @validate_sample
+    def log_prob(self, value):
+        return self.base_dist.log_prob(value) - jnp.log(self._cdf_at_high)
+
     def tree_flatten(self):
-        base_flatten, base_aux = self.base_dist.base_dist.tree_flatten()
+        base_flatten, base_aux = self.base_dist.tree_flatten()
         if isinstance(self._support.upper_bound, (int, float)):
             return base_flatten, (type(self.base_dist), base_aux, self._support.upper_bound)
         else:
-            return (base_flatten, self._support.upper_bound), (type(self.base_dist), base_aux)
+            return (base_flatten, self.high), (type(self.base_dist), base_aux)
 
     @classmethod
     def tree_unflatten(cls, aux_data, params):
@@ -1199,7 +1211,7 @@ def __init__(self, base_dist, low=0., high=1., validate_args=None):
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low), jnp.shape(high))
-        self.base_dist = base_dist
+        self.base_dist = tree_map(lambda p: promote_shapes(p, shape=batch_shape)[0], base_dist)
         self.low, = promote_shapes(low, shape=batch_shape)
         self.high, = promote_shapes(high, shape=batch_shape)
         self._support = constraints.interval(low, high)
@@ -1209,6 +1221,20 @@ def __init__(self, base_dist, low=0., high=1., validate_args=None):
     def support(self):
         return self._support
 
+    @lazy_property
+    def _tail_prob_at_low(self):
+        # if low < loc, returns cdf(low); otherwise returns 1 - cdf(low)
+        loc = self.base_dist.loc
+        sign = jnp.where(loc >= self.low, 1., -1.)
+        return self.base_dist.cdf(loc - sign * (loc - self.low))
+
+    @lazy_property
+    def _tail_prob_at_high(self):
+        # if low < loc, returns cdf(high); otherwise returns 1 - cdf(high)
+        loc = self.base_dist.loc
+        sign = jnp.where(loc >= self.low, 1., -1.)
+        return self.base_dist.cdf(loc - sign * (loc - self.high))
+
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
@@ -1220,13 +1246,10 @@ def sample(self, key, sample_shape=()):
         #   A = icdf[(1 - u) * cdf(low) + u * cdf(high)]
         # Else
         #   A = 2 * loc - icdf[(1 - u) * cdf(2*loc-low)) + u * cdf(2*loc - high)]
-        # TODO: cache sign, low_cdf, high_cdf
         loc = self.base_dist.loc
-        diff = loc - self.low
-        sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(loc - sign * diff)
-        high_cdf = self.base_dist.cdf(loc - sign * (loc - self.high))
-        return (1 - sign) * loc + sign * self.base_dist.icdf((1 - u) * low_cdf + u * high_cdf)
+        sign = jnp.where(loc >= self.low, 1., -1.)
+        return (1 - sign) * loc + sign * self.base_dist.icdf(
+            (1 - u) * self._tail_prob_at_low + u * self._tail_prob_at_high)
 
     @validate_sample
     def log_prob(self, value):
@@ -1235,12 +1258,9 @@ def log_prob(self, value):
         #   cdf(high) - cdf(low) = as-is
         # if low > loc
         #   cdf(high) - cdf(low) = cdf(2 * loc - low) - cdf(2 * loc - high)
-        loc = self.base_dist.loc
-        diff = loc - self.low
-        sign = jnp.where(diff >= 0, 1., -1.)
-        low_cdf = self.base_dist.cdf(loc - sign * diff)
-        high_cdf = self.base_dist.cdf(loc - sign * (loc - self.high))
-        return self.base_dist.log_prob(value) - jnp.log(sign * (high_cdf - low_cdf))
+        sign = jnp.where(self.base_dist.loc >= self.low, 1., -1.)
+        return self.base_dist.log_prob(value) - \
+            jnp.log(sign * (self._tail_prob_at_high - self._tail_prob_at_low))
 
     def tree_flatten(self):
         base_flatten, base_aux = self.base_dist.tree_flatten()
@@ -1249,8 +1269,7 @@ def tree_flatten(self):
             return base_flatten, (type(self.base_dist), base_aux,
                                   self._support.lower_bound, self._support.upper_bound)
         else:
-            return (base_flatten, self._support.lower_bound, self.support.upper_bound), \
-                (type(self.base_dist), base_aux)
+            return (base_flatten, self.low, self.high), (type(self.base_dist), base_aux)
 
     @classmethod
     def tree_unflatten(cls, aux_data, params):
diff --git a/test/test_distributions.py b/test/test_distributions.py
index f9f982794..3d4590c4f 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -76,6 +76,7 @@ def _TruncatedNormal(loc, scale, low, high):
     return dist.TruncatedDistribution(dist.Normal(loc, scale), low, high)
 
 
+_TruncatedNormal.arg_constraints = {}
 _TruncatedNormal.reparametrized_params = []
 _TruncatedNormal.infer_shapes = lambda *args: (lax.broadcast_shapes(*args), ())
 
@@ -204,6 +205,7 @@ def sample(self, key, sample_shape=()):
     T(_TruncatedNormal, -1., jnp.array([2., 3.]), 1., None),
     T(_TruncatedNormal, -1., 2., jnp.array([-6., 4.]), jnp.array([-4., 6.])),
     T(_TruncatedNormal, jnp.array([0., 1.]), jnp.array([[1.], [2.]]), None, jnp.array([-2., 2.])),
+    T(dist.continuous.TwoSidedTruncatedDistribution, dist.Laplace(0., 1.), -2., 3.),
     T(dist.Uniform, 0., 2.),
     T(dist.Uniform, 1., jnp.array([2., 3.])),
     T(dist.Uniform, jnp.array([0., 0.]), jnp.array([[2.], [3.]])),
@@ -372,6 +374,7 @@ def test_dist_shape(jax_dist, sp_dist, params, prepend_shape):
 @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS + DISCRETE + DIRECTIONAL)
 def test_infer_shapes(jax_dist, sp_dist, params, prepend_shape):
     shapes = tuple(getattr(p, "shape", ()) for p in params)
+    shapes = tuple(x() if callable(x) else x for x in shapes)
     try:
         expected_batch_shape, expected_event_shape = jax_dist.infer_shapes(*shapes)
     except NotImplementedError:
@@ -761,6 +764,8 @@ def fn(*args):
 
     eps = 1e-3
     for i in range(len(params)):
+        if isinstance(params[i], dist.Distribution):  # skip taking grad w.r.t. base_dist
+            continue
         if params[i] is None or jnp.result_type(params[i]) in (jnp.int32, jnp.int64):
             continue
         actual_grad = jax.grad(fn, i)(*params)
@@ -782,6 +787,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 in (_TruncatedNormal, dist.continuous.TwoSidedTruncatedDistribution):
+        pytest.skip("Truncated distributions do not has mean/var implemented")
 
     n = 20000 if jax_dist in [dist.LKJ, dist.LKJCholesky] else 200000
     d_jax = jax_dist(*params)
@@ -856,6 +863,8 @@ def test_mean_var(jax_dist, sp_dist, params):
     (2, 3),
 ])
 def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape):
+    if jax_dist is _TruncatedNormal:
+        pytest.skip("_TruncatedNormal is a function, not a class")
     dist_args = [p for p in inspect.getfullargspec(jax_dist.__init__)[0][1:]]
 
     valid_params, oob_params = list(params), list(params)
@@ -864,6 +873,8 @@ def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape):
     for i in range(len(params)):
         if jax_dist in (_ImproperWrapper, dist.LKJ, dist.LKJCholesky) and dist_args[i] != "concentration":
             continue
+        if jax_dist is dist.continuous.TwoSidedTruncatedDistribution and dist_args[i] == "base_dist":
+            continue
         if jax_dist is dist.GaussianRandomWalk and dist_args[i] == "num_steps":
             continue
         if params[i] is None:

From afc00de52683918bc9b7669f84c3cc07fac90703 Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Tue, 16 Feb 2021 23:47:41 -0600
Subject: [PATCH 6/7] address comment and add algebraic tests

---
 numpyro/distributions/continuous.py | 93 +++++++++++++++--------------
 test/test_distributions.py          | 23 ++++---
 2 files changed, 63 insertions(+), 53 deletions(-)

diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
index 26f74bb4f..acfd7112e 100644
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -677,6 +677,44 @@ def tree_flatten(self):
         return super(TransformedDistribution, self).tree_flatten()
 
 
+class Logistic(Distribution):
+    arg_constraints = {'loc': constraints.real, 'scale': constraints.positive}
+    support = constraints.real
+    reparametrized_params = ['loc', 'scale']
+
+    def __init__(self, loc=0., scale=1., validate_args=None):
+        self.loc, self.scale = promote_shapes(loc, scale)
+        batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
+        super(Logistic, self).__init__(batch_shape, validate_args=validate_args)
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        z = random.logistic(key, shape=sample_shape + self.batch_shape + self.event_shape)
+        return self.loc + z * self.scale
+
+    @validate_sample
+    def log_prob(self, value):
+        log_exponent = (self.loc - value) / self.scale
+        log_denominator = jnp.log(self.scale) + 2 * nn.softplus(log_exponent)
+        return log_exponent - log_denominator
+
+    @property
+    def mean(self):
+        return jnp.broadcast_to(self.loc, self.batch_shape)
+
+    @property
+    def variance(self):
+        var = (self.scale ** 2) * (jnp.pi ** 2) / 3
+        return jnp.broadcast_to(var, self.batch_shape)
+
+    def cdf(self, value):
+        scaled = (value - self.loc) / self.scale
+        return expit(scaled)
+
+    def icdf(self, q):
+        return self.loc + self.scale * logit(q)
+
+
 def _batch_mahalanobis(bL, bx):
     if bL.shape[:-1] == bx.shape:
         # no need to use the below optimization procedure
@@ -1090,9 +1128,10 @@ def icdf(self, q):
 class LeftTruncatedDistribution(Distribution):
     arg_constraints = {"low": constraints.real}
     reparametrized_params = ["low"]
+    supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT)
 
     def __init__(self, base_dist, low=0., validate_args=None):
-        assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
+        assert isinstance(base_dist, self.supported_types)
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low))
@@ -1115,9 +1154,7 @@ def _tail_prob_at_low(self):
     @lazy_property
     def _tail_prob_at_high(self):
         # if low < loc, returns cdf(high) = 1; otherwise returns 1 - cdf(high) = 0
-        loc = self.base_dist.loc
-        sign = jnp.where(loc >= self.low, 1., -1.)
-        return 0.5 * (1 + sign)
+        return jnp.where(self.low < self.base_dist.loc, 1., 0.)
 
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
@@ -1155,9 +1192,10 @@ def tree_unflatten(cls, aux_data, params):
 class RightTruncatedDistribution(Distribution):
     arg_constraints = {"high": constraints.real}
     reparametrized_params = ["high"]
+    supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT)
 
     def __init__(self, base_dist, high=0., validate_args=None):
-        assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
+        assert isinstance(base_dist, self.supported_types)
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(high))
@@ -1205,9 +1243,10 @@ def tree_unflatten(cls, aux_data, params):
 class TwoSidedTruncatedDistribution(Distribution):
     arg_constraints = {"low": constraints.dependent, "high": constraints.dependent}
     reparametrized_params = ["low", "high"]
+    supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT)
 
     def __init__(self, base_dist, low=0., high=1., validate_args=None):
-        assert isinstance(base_dist, (Cauchy, Laplace, Logistic, Normal, StudentT))
+        assert isinstance(base_dist, self.supported_types)
         assert base_dist.support is constraints.real, \
             "The base distribution should be univariate and have real support."
         batch_shape = lax.broadcast_shapes(base_dist.batch_shape, jnp.shape(low), jnp.shape(high))
@@ -1239,7 +1278,7 @@ def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         u = random.uniform(key, sample_shape + self.batch_shape)
 
-        # NB: we use a more numerical formula for a symmetric base distribution
+        # NB: we use a more numerically stable formula for a symmetric base distribution
         #   A = icdf(cdf(low) + (cdf(high) - cdf(low)) * u) = icdf[(1 - u) * cdf(low) + u * cdf(high)]
         # will suffer by precision issues when low is large;
         # If low < loc:
@@ -1253,7 +1292,7 @@ def sample(self, key, sample_shape=()):
 
     @validate_sample
     def log_prob(self, value):
-        # NB: we use a more numerical formula for a symmetric base distribution
+        # NB: we use a more numerically stable formula for a symmetric base distribution
         # if low < loc
         #   cdf(high) - cdf(low) = as-is
         # if low > loc
@@ -1508,44 +1547,6 @@ def infer_shapes(low=(), high=()):
         return batch_shape, event_shape
 
 
-class Logistic(Distribution):
-    arg_constraints = {'loc': constraints.real, 'scale': constraints.positive}
-    support = constraints.real
-    reparametrized_params = ['loc', 'scale']
-
-    def __init__(self, loc=0., scale=1., validate_args=None):
-        self.loc, self.scale = promote_shapes(loc, scale)
-        batch_shape = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
-        super(Logistic, self).__init__(batch_shape, validate_args=validate_args)
-
-    def sample(self, key, sample_shape=()):
-        assert is_prng_key(key)
-        z = random.logistic(key, shape=sample_shape + self.batch_shape + self.event_shape)
-        return self.loc + z * self.scale
-
-    @validate_sample
-    def log_prob(self, value):
-        log_exponent = (self.loc - value) / self.scale
-        log_denominator = jnp.log(self.scale) + 2 * nn.softplus(log_exponent)
-        return log_exponent - log_denominator
-
-    @property
-    def mean(self):
-        return jnp.broadcast_to(self.loc, self.batch_shape)
-
-    @property
-    def variance(self):
-        var = (self.scale ** 2) * (jnp.pi ** 2) / 3
-        return jnp.broadcast_to(var, self.batch_shape)
-
-    def cdf(self, value):
-        scaled = (value - self.loc) / self.scale
-        return expit(scaled)
-
-    def icdf(self, q):
-        return self.loc + self.scale * logit(q)
-
-
 class TruncatedPolyaGamma(Distribution):
     truncation_point = 2.5
     num_log_prob_terms = 7
diff --git a/test/test_distributions.py b/test/test_distributions.py
index 3d4590c4f..e43b52567 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -550,20 +550,29 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
 
 @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS)
 def test_cdf_and_icdf(jax_dist, sp_dist, params):
+    d = jax_dist(*params)
+    samples = d.sample(key=random.PRNGKey(0))
+    quantiles = random.uniform(random.PRNGKey(1), d.shape())
+    try:
+        if d.shape() == ():
+            assert_allclose(jax.grad(d.cdf)(samples), jnp.exp(d.log_prob(samples)), rtol=1e-5)
+            assert_allclose(jax.grad(d.icdf)(quantiles), jnp.exp(-d.log_prob(d.icdf(quantiles))), rtol=1e-5)
+        assert_allclose(d.cdf(d.icdf(quantiles)), quantiles, rtol=1e-5)
+        assert_allclose(d.icdf(d.cdf(samples)), samples, rtol=1e-5)
+    except NotImplementedError:
+        pass
+
+    # test against scipy
     if not sp_dist:
         pytest.skip('no corresponding scipy distn.')
-    jax_dist = jax_dist(*params)
-    if jax_dist.event_dim > 0:
+    if d.event_dim > 0:
         pytest.skip('skip testing cdf/icdf methods of multivariate distributions')
     sp_dist = sp_dist(*params)
-    rng_key = random.PRNGKey(0)
-    samples = jax_dist.sample(key=rng_key)
-    quantiles = random.uniform(random.PRNGKey(1), jax_dist.shape())
     try:
-        actual_cdf = jax_dist.cdf(samples)
+        actual_cdf = d.cdf(samples)
         expected_cdf = sp_dist.cdf(samples)
         assert_allclose(actual_cdf, expected_cdf, rtol=1e-5)
-        actual_icdf = jax_dist.icdf(quantiles)
+        actual_icdf = d.icdf(quantiles)
         expected_icdf = sp_dist.ppf(quantiles)
         assert_allclose(actual_icdf, expected_icdf, rtol=1e-5)
     except NotImplementedError:

From db47ab38faef88630ba9a8b2586132e4ba881b97 Mon Sep 17 00:00:00 2001
From: Du Phan <fehiepsi@gmail.com>
Date: Wed, 17 Feb 2021 00:34:10 -0600
Subject: [PATCH 7/7] test cdf/icdf with more random values

---
 test/test_distributions.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/test/test_distributions.py b/test/test_distributions.py
index e43b52567..76602bd5f 100644
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -551,30 +551,33 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
 @pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS)
 def test_cdf_and_icdf(jax_dist, sp_dist, params):
     d = jax_dist(*params)
-    samples = d.sample(key=random.PRNGKey(0))
-    quantiles = random.uniform(random.PRNGKey(1), d.shape())
+    if d.event_dim > 0:
+        pytest.skip('skip testing cdf/icdf methods of multivariate distributions')
+    samples = d.sample(key=random.PRNGKey(0), sample_shape=(100,))
+    quantiles = random.uniform(random.PRNGKey(1), (100,) + d.shape())
     try:
         if d.shape() == ():
-            assert_allclose(jax.grad(d.cdf)(samples), jnp.exp(d.log_prob(samples)), rtol=1e-5)
-            assert_allclose(jax.grad(d.icdf)(quantiles), jnp.exp(-d.log_prob(d.icdf(quantiles))), rtol=1e-5)
-        assert_allclose(d.cdf(d.icdf(quantiles)), quantiles, rtol=1e-5)
-        assert_allclose(d.icdf(d.cdf(samples)), samples, rtol=1e-5)
+            rtol = 1e-3 if jax_dist is dist.StudentT else 1e-5
+            assert_allclose(jax.vmap(jax.grad(d.cdf))(samples),
+                            jnp.exp(d.log_prob(samples)), atol=1e-5, rtol=rtol)
+            assert_allclose(jax.vmap(jax.grad(d.icdf))(quantiles),
+                            jnp.exp(-d.log_prob(d.icdf(quantiles))), atol=1e-5, rtol=rtol)
+        assert_allclose(d.cdf(d.icdf(quantiles)), quantiles, atol=1e-5, rtol=1e-5)
+        assert_allclose(d.icdf(d.cdf(samples)), samples, atol=1e-5, rtol=1e-5)
     except NotImplementedError:
         pass
 
     # test against scipy
     if not sp_dist:
         pytest.skip('no corresponding scipy distn.')
-    if d.event_dim > 0:
-        pytest.skip('skip testing cdf/icdf methods of multivariate distributions')
     sp_dist = sp_dist(*params)
     try:
         actual_cdf = d.cdf(samples)
         expected_cdf = sp_dist.cdf(samples)
-        assert_allclose(actual_cdf, expected_cdf, rtol=1e-5)
+        assert_allclose(actual_cdf, expected_cdf, atol=1e-5, rtol=1e-5)
         actual_icdf = d.icdf(quantiles)
         expected_icdf = sp_dist.ppf(quantiles)
-        assert_allclose(actual_icdf, expected_icdf, rtol=1e-5)
+        assert_allclose(actual_icdf, expected_icdf, atol=1e-5, rtol=1e-4)
     except NotImplementedError:
         pass