Restricted icdf function to [0,1]

numpyro/distributions/truncated.py

@@ -72,9 +72,10 @@ def sample(self, key, sample_shape=()):
     def icdf(self, q):
         loc = self.base_dist.loc
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
-        return (1 - sign) * loc + sign * self.base_dist.icdf(
+        ppf = (1 - sign) * loc + sign * self.base_dist.icdf(
             (1 - q) * self._tail_prob_at_low + q * self._tail_prob_at_high
         )
+        return jnp.where(q < 0, jnp.nan, ppf)
 
     @validate_sample
     def log_prob(self, value):
@@ -144,7 +145,8 @@ def sample(self, key, sample_shape=()):
         return self.icdf(u)
 
     def icdf(self, q):
-        return self.base_dist.icdf(q * self._cdf_at_high)
+        ppf = self.base_dist.icdf(q * self._cdf_at_high)
+        return jnp.where(q > 1, jnp.nan, ppf)
 
     @validate_sample
     def log_prob(self, value):
@@ -253,9 +255,10 @@ def icdf(self, q):
         #   A = 2 * loc - icdf[(1 - q) * cdf(2*loc-low)) + q * cdf(2*loc - high)]
         loc = self.base_dist.loc
         sign = jnp.where(loc >= self.low, 1.0, -1.0)
-        return (1 - sign) * loc + sign * self.base_dist.icdf(
+        ppf = (1 - sign) * loc + sign * self.base_dist.icdf(
             clamp_probs((1 - q) * self._tail_prob_at_low + q * self._tail_prob_at_high)
         )
+        return jnp.where(jnp.logical_or(q < 0, q > 1), jnp.nan, ppf)
 
     @validate_sample
     def log_prob(self, value):