Repository: pyro-ppl/numpyro

Issue:
Support for TruncatedDistribution
Currently, we have one-side truncated normal and cauchy implemented. But it would be nice to add a general class, which uses `cdf` method, for other distributions. There is a [nice discussion here](https://github.com/probtorch/pytorch/pull/121) for context on relevant numerical issues.

cc @fritzo 


Discussion:
Note we also have the one-sided general [FoldedDistribution](http://docs.pyro.ai/en/stable/distributions.html#pyro.distributions.FoldedDistribution) in Pyro, which is a cheap rsample-friendly qualitatively similar alternative to `TruncatedDistribution`. This doesn't get us any closer to literally translating Stan models, but maybe we can take license to approximately translate models and e.g. check that we're only slightly truncating, or not truncating past the median (since `Folded` and `Truncated` differ for severely truncated distributions).

Resolving patch:
diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
--- 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
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -26,12 +26,12 @@
 # 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
 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 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}
@@ -109,6 +112,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 +422,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 * jnp.sign(a) * jnp.log1p(-2 * jnp.abs(a))
+
 
 class LKJ(TransformedDistribution):
     r"""
@@ -659,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
@@ -966,6 +1022,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 +1109,241 @@ def variance(self):
         var = jnp.where(self.df <= 1, jnp.nan, var)
         return jnp.broadcast_to(var, self.batch_shape)
 
+    def cdf(self, value):
+        # 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):
+        # scipy.special.betaincinv is not avaiable yet in JAX
+        # upstream issue: https://github.com/google/jax/issues/2399
+        raise NotImplementedError
+
+
+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, 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))
+        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)
+
+    @constraints.dependent_property(is_discrete=False, event_dim=0)
+    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
+        return jnp.where(self.low < self.base_dist.loc, 1., 0.)
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        u = random.uniform(key, sample_shape + self.batch_shape)
+        loc = self.base_dist.loc
+        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):
+        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.low), (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 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, 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))
+        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.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.high), (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"]
+    supported_types = (Cauchy, Laplace, Logistic, Normal, StudentT)
+
+    def __init__(self, base_dist, low=0., high=1., validate_args=None):
+        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))
+        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)
+        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 _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)
+
+        # 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:
+        #   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)]
+        loc = self.base_dist.loc
+        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 numerically stable 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)
+        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)) 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.low, self.high), (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=None, high=None, validate_args=None):
+    """
+    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:
+        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)
+
 
 class _BaseTruncatedCauchy(Distribution):
     # NB: this is a truncated cauchy with low=0, scale=1
@@ -1252,37 +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)
-
-
 class TruncatedPolyaGamma(Distribution):
     truncation_point = 2.5
     num_log_prob_terms = 7
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- 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 = {}


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -60,6 +60,27 @@ 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.arg_constraints = {}
+_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 +124,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 +200,12 @@ 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.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.]])),
@@ -346,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:
@@ -519,6 +548,40 @@ 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):
+    d = jax_dist(*params)
+    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() == ():
+            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.')
+    sp_dist = sp_dist(*params)
+    try:
+        actual_cdf = d.cdf(samples)
+        expected_cdf = sp_dist.cdf(samples)
+        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, atol=1e-5, rtol=1e-4)
+    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__:
@@ -713,6 +776,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)
@@ -734,6 +799,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)
@@ -808,6 +875,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)
@@ -816,6 +885,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: