Generalize handling of invalid values in measurable transforms

We rely on the jacobian returning `nan` for invalid values

Co-authored-by: Luke LB <[email protected]>

Author: ricardoV94

pymc/logprob/transforms.py

@@ -374,9 +374,14 @@ def measurable_transform_logprob(op: MeasurableTransform, values, *inputs, **kwa
     else:
         input_logprob = logprob(measurable_input, backward_value)
 
+    if input_logprob.ndim < value.ndim:
+        # Do we just need to sum the jacobian terms across the support dims?
+        raise NotImplementedError("Transform of multivariate RVs not implemented")
+
     jacobian = op.transform_elemwise.log_jac_det(value, *other_inputs)
 
-    return input_logprob + jacobian
+    # The jacobian is used to ensure a value in the supported domain was provided
+    return at.switch(at.isnan(jacobian), -np.inf, input_logprob + jacobian)
 
 
 @node_rewriter([reciprocal])
@@ -711,18 +716,32 @@ def forward(self, value, *inputs):
         at.power(value, self.power)
 
     def backward(self, value, *inputs):
-        backward_value = at.power(value, (1 / self.power))
+        inv_power = 1 / self.power
+
+        # Powers that don't admit negative values
+        if (np.abs(self.power) < 1) or (self.power % 2 == 0):
+            backward_value = at.switch(value >= 0, at.power(value, inv_power), np.nan)
+        # Powers that admit negative values require special logic, because (-1)**(1/3) returns `nan` in PyTensor
+        else:
+            backward_value = at.power(at.abs(value), inv_power) * at.switch(value >= 0, 1, -1)
 
         # In this case the transform is not 1-to-1
-        if (self.power > 1) and (self.power % 2 == 0):
+        if self.power % 2 == 0:
             return -backward_value, backward_value
         else:
             return backward_value
 
     def log_jac_det(self, value, *inputs):
         inv_power = 1 / self.power
+
         # Note: This fails for value==0
-        return np.log(np.abs(inv_power)) + (inv_power - 1) * at.log(value)
+        res = np.log(np.abs(inv_power)) + (inv_power - 1) * at.log(at.abs(value))
+
+        # Powers that don't admit negative values
+        if (np.abs(self.power) < 1) or (self.power % 2 == 0):
+            res = at.switch(value >= 0, res, np.nan)
+
+        return res
 
 
 class IntervalTransform(RVTransform):

pymc/tests/logprob/test_transforms.py

@@ -632,15 +632,15 @@ def test_chained_transform():
 
 
 def test_exp_transform_rv():
-    base_rv = at.random.normal(0, 1, size=2, name="base_rv")
+    base_rv = at.random.normal(0, 1, size=3, name="base_rv")
     y_rv = at.exp(base_rv)
     y_rv.name = "y"
 
     y_vv = y_rv.clone()
     logp = joint_logprob({y_rv: y_vv}, sum=False)
     logp_fn = pytensor.function([y_vv], logp)
 
-    y_val = [0.1, 0.3]
+    y_val = [-2.0, 0.1, 0.3]
     np.testing.assert_allclose(
         logp_fn(y_val),
         sp.stats.lognorm(s=1).logpdf(y_val),
@@ -794,28 +794,28 @@ def test_invalid_broadcasted_transform_rv_fails():
 def test_reciprocal_rv_transform(numerator):
     shape = 3
     scale = 5
-    x_rv = numerator / at.random.gamma(shape, scale)
+    x_rv = numerator / at.random.gamma(shape, scale, size=(2,))
     x_rv.name = "x"
 
     x_vv = x_rv.clone()
-    x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}))
+    x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}, sum=False))
 
-    x_test_val = 1.5
-    assert np.isclose(
+    x_test_val = np.r_[-0.5, 1.5]
+    assert np.allclose(
         x_logp_fn(x_test_val),
         sp.stats.invgamma(shape, scale=scale * numerator).logpdf(x_test_val),
     )
 
 
 def test_sqr_transform():
     # The square of a unit normal is a chi-square with 1 df
-    x_rv = at.random.normal(0, 1, size=(3,)) ** 2
+    x_rv = at.random.normal(0, 1, size=(4,)) ** 2
     x_rv.name = "x"
 
     x_vv = x_rv.clone()
     x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}, sum=False))
 
-    x_test_val = np.r_[0.5, 1, 2.5]
+    x_test_val = np.r_[-0.5, 0.5, 1, 2.5]
     assert np.allclose(
         x_logp_fn(x_test_val),
         sp.stats.chi2(df=1).logpdf(x_test_val),
@@ -824,19 +824,58 @@ def test_sqr_transform():
 
 def test_sqrt_transform():
     # The sqrt of a chisquare with n df is a chi distribution with n df
-    x_rv = at.sqrt(at.random.chisquare(df=3, size=(3,)))
+    x_rv = at.sqrt(at.random.chisquare(df=3, size=(4,)))
     x_rv.name = "x"
 
     x_vv = x_rv.clone()
     x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}, sum=False))
 
-    x_test_val = np.r_[0.5, 1, 2.5]
+    x_test_val = np.r_[-2.5, 0.5, 1, 2.5]
     assert np.allclose(
         x_logp_fn(x_test_val),
         sp.stats.chi(df=3).logpdf(x_test_val),
     )
 
 
+@pytest.mark.parametrize("power", (-3, -1, 1, 5, 7))
+def test_negative_value_odd_power_transform(power):
+    # check that negative values and odd powers evaluate to a finite logp
+    x_rv = at.random.normal() ** power
+    x_rv.name = "x"
+
+    x_vv = x_rv.clone()
+    x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}, sum=False))
+
+    assert np.isfinite(x_logp_fn(1))
+    assert np.isfinite(x_logp_fn(-1))
+
+
+@pytest.mark.parametrize("power", (-2, 2, 4, 6, 8))
+def test_negative_value_even_power_transform(power):
+    # check that negative values and odd powers evaluate to -inf logp
+    x_rv = at.random.normal() ** power
+    x_rv.name = "x"
+
+    x_vv = x_rv.clone()
+    x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}, sum=False))
+
+    assert np.isfinite(x_logp_fn(1))
+    assert np.isneginf(x_logp_fn(-1))
+
+
+@pytest.mark.parametrize("power", (-1 / 3, -1 / 2, 1 / 2, 1 / 3))
+def test_negative_value_frac_power_transform(power):
+    # check that negative values and fractional powers evaluate to -inf logp
+    x_rv = at.random.normal() ** power
+    x_rv.name = "x"
+
+    x_vv = x_rv.clone()
+    x_logp_fn = pytensor.function([x_vv], joint_logprob({x_rv: x_vv}, sum=False))
+
+    assert np.isfinite(x_logp_fn(2.5))
+    assert np.isneginf(x_logp_fn(-2.5))
+
+
 def test_negated_rv_transform():
     x_rv = -at.random.halfnormal()
     x_rv.name = "x"