Repository: pyro-ppl/numpyro

Issue:
Allow for `high` in TruncatedNormal, TruncatedCauchy
TruncatedDistribution has both `low` and `high`. Why do `TruncatedNormal` and `TruncatedCauchy` only have `low`?


Discussion:
Hi @MarcoGorelli, the reason is currently, the signatures of `TruncatedNormal` is `low, loc, scale`, so introducing `low, high, loc, scale` will break current user codes. I think we can introduce the `high` parameter but forcing parameter names to be explicit (at least for the next release) so that users can be informed of affected codes
```python
def TruncatedNormal(*, loc=0., scale=1., low=None, high=None, validate_args=None):
    base_dist = dist.Normal(loc, scale)
    return TruncatedDistribution(base_dist, low=None, high=None, validate_args=None)

def TruncatedCauchy(*, loc=0., scale=1., low=None, high=None, validate_args=None):
    base_dist = dist.Normal(loc, scale)
    return TruncatedDistribution(base_dist, low=None, high=None, validate_args=None)
```

To get the mean/variance of TruncatedNormal/TruncatedCauchy, I think you can add `mean`/`variance` methods to `LeftTruncatedDistribution` (and probably also for `RightTruncatedDistribution`, `TwoSidedTruncatedDistribution` if you come up with some formula), kind of
```
class LeftTruncatedDistribution(...):
    ...
    @property
    def mean(self):
         if isinstance(self.base_dist, dist.Normal):
              # implement
         elif isinstance(self.base_dist, dist.Cauchy):
              # ...
         else:
              raise NotImplementedError
```
Thanks @fehiepsi !

So a warning could be added for 0.7.3 which warns that all args of TruncatdNormal will become keyword-only, and then in 0.8.0 the change is enforced?

> To get the mean/variance of TruncatedNormal/TruncatedCauchy,

Sorry, I don't understand - why does this need adding in order to have `high` in `TruncatedNormal`? Or is it a separate enhancement?
I guess no need for a warning because Python will inform those argument changes when non-keyword arguments are provided. Regarding mean/variance, currently the LeftTruncated or RightTruncated does not provide mean/variance properties, but we have them in TruncatedNormal and TruncatedCauchy. The code in my last comment is just to make sure that the change does not affect those users who are using TruncatedNormal.mean.
Sorry for the delay, I'll look into this this week

Resolving patch:
diff --git a/numpyro/contrib/tfp/distributions.py b/numpyro/contrib/tfp/distributions.py
--- a/numpyro/contrib/tfp/distributions.py
+++ b/numpyro/contrib/tfp/distributions.py
@@ -163,6 +163,13 @@ def init(self, *args, **kwargs):
                 "loc": constraints.real,
                 "scale": constraints.positive,
             }
+        elif tfd_class is tfd.TruncatedCauchy:
+            _PyroDist.arg_constraints = {
+                "low": constraints.real,
+                "high": constraints.real,
+                "loc": constraints.real,
+                "scale": constraints.positive,
+            }
         else:
             if hasattr(numpyro_dist, tfd_class_name):
                 numpyro_dist_class = getattr(numpyro_dist, tfd_class_name)
diff --git a/numpyro/distributions/truncated.py b/numpyro/distributions/truncated.py
--- a/numpyro/distributions/truncated.py
+++ b/numpyro/distributions/truncated.py
@@ -105,8 +105,10 @@ def mean(self):
         if isinstance(self.base_dist, Normal):
             low_prob = jnp.exp(self.log_prob(self.low))
             return self.base_dist.loc + low_prob * self.base_dist.scale ** 2
+        elif isinstance(self.base_dist, Cauchy):
+            return jnp.full(self.batch_shape, jnp.nan)
         else:
-            return NotImplementedError("mean only available for Normal")
+            return NotImplementedError("mean only available for Normal and Cauchy")
 
     @property
     def var(self):
@@ -117,8 +119,10 @@ def var(self):
                 + (self.low - self.base_dist.loc) * low_prob
                 - (low_prob * self.base_dist.scale) ** 2
             )
+        elif isinstance(self.base_dist, Cauchy):
+            return jnp.full(self.batch_shape, jnp.nan)
         else:
-            return NotImplementedError("var only available for Normal")
+            return NotImplementedError("var only available for Normal and Cauchy")
 
 
 class RightTruncatedDistribution(Distribution):
@@ -186,8 +190,10 @@ def mean(self):
         if isinstance(self.base_dist, Normal):
             high_prob = jnp.exp(self.log_prob(self.high))
             return self.base_dist.loc - high_prob * self.base_dist.scale ** 2
+        elif isinstance(self.base_dist, Cauchy):
+            return jnp.full(self.batch_shape, jnp.nan)
         else:
-            return NotImplementedError("mean only available for Normal")
+            return NotImplementedError("mean only available for Normal and Cauchy")
 
     @property
     def var(self):
@@ -198,8 +204,10 @@ def var(self):
                 - (self.high - self.base_dist.loc) * high_prob
                 - (high_prob * self.base_dist.scale) ** 2
             )
+        elif isinstance(self.base_dist, Cauchy):
+            return jnp.full(self.batch_shape, jnp.nan)
         else:
-            return NotImplementedError("var only available for Normal")
+            return NotImplementedError("var only available for Normal and Cauchy")
 
 
 class TwoSidedTruncatedDistribution(Distribution):
@@ -306,8 +314,10 @@ def mean(self):
             return (
                 self.base_dist.loc + (low_prob - high_prob) * self.base_dist.scale ** 2
             )
+        elif isinstance(self.base_dist, Cauchy):
+            return jnp.full(self.batch_shape, jnp.nan)
         else:
-            return NotImplementedError("mean only available for Normal")
+            return NotImplementedError("mean only available for Normal and Cauchy")
 
     @property
     def var(self):
@@ -320,8 +330,10 @@ def var(self):
                 - (self.high - self.base_dist.loc) * high_prob
                 - ((low_prob - high_prob) * self.base_dist.scale) ** 2
             )
+        elif isinstance(self.base_dist, Cauchy):
+            return jnp.full(self.batch_shape, jnp.nan)
         else:
-            return NotImplementedError("var only available for Normal")
+            return NotImplementedError("var only available for Normal and Cauchy")
 
 
 def TruncatedDistribution(base_dist, low=None, high=None, validate_args=None):
@@ -353,44 +365,13 @@ def TruncatedDistribution(base_dist, low=None, high=None, validate_args=None):
         )
 
 
-class TruncatedCauchy(LeftTruncatedDistribution):
-    arg_constraints = {
-        "low": constraints.real,
-        "loc": constraints.real,
-        "scale": constraints.positive,
-    }
-    reparametrized_params = ["low", "loc", "scale"]
-
-    def __init__(self, low=0.0, loc=0.0, scale=1.0, validate_args=None):
-        self.low, self.loc, self.scale = promote_shapes(low, loc, scale)
-        super().__init__(
-            Cauchy(self.loc, self.scale), low=self.low, validate_args=validate_args
-        )
-
-    @property
-    def mean(self):
-        return jnp.full(self.batch_shape, jnp.nan)
-
-    @property
-    def variance(self):
-        return jnp.full(self.batch_shape, jnp.nan)
-
-    def tree_flatten(self):
-        if isinstance(self._support.lower_bound, (int, float)):
-            aux_data = self._support.lower_bound
-        else:
-            aux_data = None
-        return (self.low, self.loc, self.scale), aux_data
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, params):
-        d = cls(*params)
-        if aux_data is not None:
-            d._support = constraints.greater_than(aux_data)
-        return d
+def TruncatedCauchy(loc=0.0, scale=1.0, *, low=None, high=None, validate_args=None):
+    return TruncatedDistribution(
+        Cauchy(loc, scale), low=low, high=high, validate_args=validate_args
+    )
 
 
-def TruncatedNormal(*, loc=0.0, scale=1.0, low=None, high=None, validate_args=None):
+def TruncatedNormal(loc=0.0, scale=1.0, *, low=None, high=None, validate_args=None):
     return TruncatedDistribution(
         Normal(loc, scale), low=low, high=high, validate_args=validate_args
     )


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -101,6 +101,10 @@ def _TruncatedNormal(loc, scale, low, high):
     return dist.TruncatedNormal(loc=loc, scale=scale, low=low, high=high)
 
 
+def _TruncatedCauchy(loc, scale, low, high):
+    return dist.TruncatedCauchy(loc=loc, scale=scale, low=low, high=high)
+
+
 _TruncatedNormal.arg_constraints = {}
 _TruncatedNormal.reparametrized_params = []
 _TruncatedNormal.infer_shapes = lambda *args: (lax.broadcast_shapes(*args), ())
@@ -417,14 +421,17 @@ def get_sp_dist(jax_dist):
     T(dist.StudentT, 1.0, 1.0, 0.5),
     T(dist.StudentT, 2.0, np.array([1.0, 2.0]), 2.0),
     T(dist.StudentT, np.array([3.0, 5.0]), np.array([[1.0], [2.0]]), 2.0),
-    T(dist.TruncatedCauchy, -1.0, 0.0, 1.0),
-    T(dist.TruncatedCauchy, 1.0, 0.0, np.array([1.0, 2.0])),
+    T(_TruncatedCauchy, 0.0, 1.0, -1.0, None),
+    T(_TruncatedCauchy, 0.0, np.array([1.0, 2.0]), 1.0, None),
     T(
-        dist.TruncatedCauchy,
-        np.array([-2.0, 2.0]),
+        _TruncatedCauchy,
         np.array([0.0, 1.0]),
         np.array([[1.0], [2.0]]),
+        np.array([-2.0, 2.0]),
+        None,
     ),
+    T(_TruncatedCauchy, 0.0, 1.0, None, 1.0),
+    T(_TruncatedCauchy, 0.0, 1.0, -1.0, 1.0),
     T(_TruncatedNormal, 0.0, 1.0, -1.0, None),
     T(_TruncatedNormal, -1.0, np.array([1.0, 2.0]), 1.0, None),
     T(
@@ -984,13 +991,12 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
                     params[1],
                     params[2],
                 )
-                if low is None:
-                    low = -np.inf
-                if high is None:
-                    high = np.inf
             else:
                 # old api
-                low, loc, scale = params
+                loc, scale, low, high = params
+            if low is None:
+                low = -np.inf
+            if high is None:
                 high = np.inf
             sp_dist = get_sp_dist(type(jax_dist.base_dist))(loc, scale)
             expected = sp_dist.logpdf(samples) - jnp.log(
@@ -1332,6 +1338,7 @@ def test_mean_var(jax_dist, sp_dist, params):
         pytest.skip("Skewed Distribution are not symmetric about location.")
     if jax_dist in (
         _TruncatedNormal,
+        _TruncatedCauchy,
         dist.LeftTruncatedDistribution,
         dist.RightTruncatedDistribution,
         dist.TwoSidedTruncatedDistribution,
@@ -1437,7 +1444,12 @@ def test_mean_var(jax_dist, sp_dist, params):
 )
 @pytest.mark.parametrize("prepend_shape", [(), (2,), (2, 3)])
 def test_distribution_constraints(jax_dist, sp_dist, params, prepend_shape):
-    if jax_dist in (_TruncatedNormal, _GaussianMixture, _Gaussian2DMixture):
+    if jax_dist in (
+        _TruncatedNormal,
+        _TruncatedCauchy,
+        _GaussianMixture,
+        _Gaussian2DMixture,
+    ):
         pytest.skip(f"{jax_dist.__name__} is a function, not a class")
     dist_args = [p for p in inspect.getfullargspec(jax_dist.__init__)[0][1:]]