Allow IntervalTransform to handle dynamic infinite bounds

Fixes bug in partially observed truncated variable

Author: ricardoV94

pymc/distributions/continuous.py

@@ -664,7 +664,7 @@ class TruncatedNormal(BoundedContinuous):
     @classmethod
     def dist(
         cls,
-        mu: Optional[DIST_PARAMETER_TYPES] = None,
+        mu: Optional[DIST_PARAMETER_TYPES] = 0,
         sigma: Optional[DIST_PARAMETER_TYPES] = None,
         *,
         tau: Optional[DIST_PARAMETER_TYPES] = None,
@@ -674,7 +674,6 @@ def dist(
     ) -> RandomVariable:
         tau, sigma = get_tau_sigma(tau=tau, sigma=sigma)
         sigma = pt.as_tensor_variable(sigma)
-        tau = pt.as_tensor_variable(tau)
         mu = pt.as_tensor_variable(floatX(mu))
 
         lower = pt.as_tensor_variable(floatX(lower)) if lower is not None else pt.constant(-np.inf)

pymc/logprob/transforms.py

@@ -816,41 +816,97 @@ def __init__(self, args_fn: Callable[..., Tuple[Optional[Variable], Optional[Var
         """
         self.args_fn = args_fn
 
-    def forward(self, value, *inputs):
+    def get_a_and_b(self, inputs):
+        """Return interval bound values.
+
+        Also returns two boolean variables indicating whether the transform is known to be statically bounded.
+        This is used to generate smaller graphs in the transform methods.
+        """
         a, b = self.args_fn(*inputs)
+        lower_bounded, upper_bounded = True, True
+        if a is None:
+            a = -pt.inf
+            lower_bounded = False
+        if b is None:
+            b = pt.inf
+            upper_bounded = False
+        return a, b, lower_bounded, upper_bounded
 
-        if a is not None and b is not None:
-            return pt.log(value - a) - pt.log(b - value)
-        elif a is not None:
-            return pt.log(value - a)
-        elif b is not None:
-            return pt.log(b - value)
+    def forward(self, value, *inputs):
+        a, b, lower_bounded, upper_bounded = self.get_a_and_b(inputs)
+
+        log_lower_distance = pt.log(value - a)
+        log_upper_distance = pt.log(b - value)
+
+        if lower_bounded and upper_bounded:
+            return pt.where(
+                pt.and_(pt.neq(a, -pt.inf), pt.neq(b, pt.inf)),
+                log_lower_distance - log_upper_distance,
+                pt.where(
+                    pt.neq(a, -pt.inf),
+                    log_lower_distance,
+                    pt.where(
+                        pt.neq(b, pt.inf),
+                        log_upper_distance,
+                        value,
+                    ),
+                ),
+            )
+        elif lower_bounded:
+            return log_lower_distance
+        elif upper_bounded:
+            return log_upper_distance
         else:
-            raise ValueError("Both edges of IntervalTransform cannot be None")
+            return value
 
     def backward(self, value, *inputs):
-        a, b = self.args_fn(*inputs)
-
-        if a is not None and b is not None:
-            sigmoid_x = pt.sigmoid(value)
-            return sigmoid_x * b + (1 - sigmoid_x) * a
-        elif a is not None:
-            return pt.exp(value) + a
-        elif b is not None:
-            return b - pt.exp(value)
+        a, b, lower_bounded, upper_bounded = self.get_a_and_b(inputs)
+
+        exp_value = pt.exp(value)
+        sigmoid_x = pt.sigmoid(value)
+        lower_distance = exp_value + a
+        upper_distance = b - exp_value
+
+        if lower_bounded and upper_bounded:
+            return pt.where(
+                pt.and_(pt.neq(a, -pt.inf), pt.neq(b, pt.inf)),
+                sigmoid_x * b + (1 - sigmoid_x) * a,
+                pt.where(
+                    pt.neq(a, -pt.inf),
+                    lower_distance,
+                    pt.where(
+                        pt.neq(b, pt.inf),
+                        upper_distance,
+                        value,
+                    ),
+                ),
+            )
+        elif lower_bounded:
+            return lower_distance
+        elif upper_bounded:
+            return upper_distance
         else:
-            raise ValueError("Both edges of IntervalTransform cannot be None")
+            return value
 
     def log_jac_det(self, value, *inputs):
-        a, b = self.args_fn(*inputs)
+        a, b, lower_bounded, upper_bounded = self.get_a_and_b(inputs)
 
-        if a is not None and b is not None:
+        if lower_bounded and upper_bounded:
             s = pt.softplus(-value)
-            return pt.log(b - a) - 2 * s - value
-        elif a is None and b is None:
-            raise ValueError("Both edges of IntervalTransform cannot be None")
-        else:
+
+            return pt.where(
+                pt.and_(pt.neq(a, -pt.inf), pt.neq(b, pt.inf)),
+                pt.log(b - a) - 2 * s - value,
+                pt.where(
+                    pt.or_(pt.neq(a, -pt.inf), pt.neq(b, pt.inf)),
+                    value,
+                    pt.zeros_like(value),
+                ),
+            )
+        elif lower_bounded or upper_bounded:
             return value
+        else:
+            return pt.zeros_like(value)
 
 
 class LogOddsTransform(RVTransform):

tests/distributions/test_truncated.py

@@ -20,7 +20,12 @@
 from pytensor.tensor.random.basic import GeometricRV, NormalRV
 
 from pymc import Censored, Model, draw, find_MAP
-from pymc.distributions.continuous import Exponential, Gamma, TruncatedNormalRV
+from pymc.distributions.continuous import (
+    Exponential,
+    Gamma,
+    TruncatedNormal,
+    TruncatedNormalRV,
+)
 from pymc.distributions.shape_utils import change_dist_size
 from pymc.distributions.transforms import _default_transform
 from pymc.distributions.truncated import Truncated, TruncatedRV, _truncated
@@ -423,3 +428,53 @@ def test_truncated_gamma():
         logp_resized_pymc,
         logp_scipy,
     )
+
+
+def test_vectorized_bounds():
+    with Model() as m:
+        x1 = TruncatedNormal("x1", lower=None, upper=0, initval=-1)
+        x2 = TruncatedNormal("x2", lower=0, upper=None, initval=1)
+        x3 = TruncatedNormal("x3", lower=-np.pi, upper=np.e, initval=-1)
+        x4 = TruncatedNormal("x4", lower=None, upper=None, initval=1)
+
+        xs = TruncatedNormal(
+            "xs",
+            lower=[-np.inf, 0, -np.pi, -np.inf],
+            upper=[0, np.inf, np.e, np.inf],
+            initval=[-1, 1, -1, 1],
+        )
+        xs_sym = Truncated(
+            "xs_sym",
+            dist=rejection_normal(),
+            lower=[-np.inf, 0, -np.pi, -np.inf],
+            upper=[0, np.inf, np.e, np.inf],
+            initval=[-1, 1, -1, 1],
+        )
+
+    ip = m.initial_point()
+    np.testing.assert_allclose(
+        np.stack([ip[f"x{i + 1}_interval__"] for i in range(4)]),
+        ip["xs_interval__"],
+    )
+    np.testing.assert_allclose(
+        ip["xs_interval__"],
+        ip["xs_sym_interval__"],
+    )
+    np.testing.assert_allclose(
+        m.rvs_to_transforms[xs].backward(ip["xs_interval__"], *xs.owner.inputs).eval(),
+        [-1, 1, -1, 1],
+    )
+    np.testing.assert_allclose(
+        m.rvs_to_transforms[xs_sym].backward(ip["xs_sym_interval__"], *xs_sym.owner.inputs).eval(),
+        [-1, 1, -1, 1],
+    )
+    *x_logp, xs_logp, xs_sym_logp = m.compile_logp(sum=False)(ip)
+    assert np.all(np.isfinite(xs_logp))
+    np.testing.assert_allclose(
+        np.stack(x_logp),
+        xs_logp,
+    )
+    np.testing.assert_allclose(
+        xs_logp,
+        xs_sym_logp,
+    )

tests/logprob/test_transforms.py

@@ -55,7 +55,6 @@
     ErfcxTransform,
     ErfTransform,
     ExpTransform,
-    IntervalTransform,
     LocTransform,
     LogTransform,
     RVTransform,
@@ -139,23 +138,6 @@ def backward(self, value, *inputs):
         ):
             SquareTransform().log_jac_det(0)
 
-    def test_invalid_interval_transform(self):
-        x_rv = pt.random.normal(0, 1)
-        x_vv = x_rv.clone()
-
-        msg = "Both edges of IntervalTransform cannot be None"
-        tr = IntervalTransform(lambda *inputs: (None, None))
-        with pytest.raises(ValueError, match=msg):
-            tr.forward(x_vv, *x_rv.owner.inputs)
-
-        tr = IntervalTransform(lambda *inputs: (None, None))
-        with pytest.raises(ValueError, match=msg):
-            tr.backward(x_vv, *x_rv.owner.inputs)
-
-        tr = IntervalTransform(lambda *inputs: (None, None))
-        with pytest.raises(ValueError, match=msg):
-            tr.log_jac_det(x_vv, *x_rv.owner.inputs)
-
     def test_chained_transform(self):
         loc = 5
         scale = 0.1

tests/model/test_core.py

@@ -40,7 +40,7 @@
 
 import pymc as pm
 
-from pymc import Deterministic, Potential
+from pymc import Deterministic, Model, Potential
 from pymc.blocking import DictToArrayBijection, RaveledVars
 from pymc.distributions import Normal, transforms
 from pymc.distributions.distribution import PartialObservedRV
@@ -1524,6 +1524,18 @@ def test_symbolic_random_variable(self):
             st.norm.logcdf(0) * 10,
         )
 
+    def test_truncated_normal(self):
+        """Test transform of unobserved TruncatedNormal leads to finite logp.
+
+        Regression test for #6999
+        """
+        with Model() as m:
+            mu = pm.TruncatedNormal("mu", mu=1, sigma=2, lower=0)
+            x = pm.TruncatedNormal(
+                "x", mu=mu, sigma=0.5, lower=0, observed=np.array([0.1, 0.2, 0.5, np.nan, np.nan])
+            )
+        m.check_start_vals(m.initial_point())
+
 
 class TestShared:
     def test_deterministic(self):