Fix probability of cosh transform

Author: ricardoV94

pymc/logprob/transforms.py

@@ -679,7 +679,15 @@ def forward(self, value, *inputs):
         return pt.cosh(value)
 
     def backward(self, value, *inputs):
-        return pt.arccosh(value)
+        back_value = pt.arccosh(value)
+        return (-back_value, back_value)
+
+    def log_jac_det(self, value, *inputs):
+        return pt.switch(
+            value < 1,
+            np.nan,
+            -pt.log(pt.sqrt(value**2 - 1)),
+        )
 
 
 class TanhTransform(RVTransform):

tests/distributions/test_transform.py

@@ -97,6 +97,9 @@ def check_jacobian_det(
     rv_inputs = rv_var.owner.inputs if rv_var.owner else []
 
     x = transform.backward(y, *rv_inputs)
+    # Assume non-injective transforms are symmetric around the origin
+    if isinstance(x, tuple):
+        x = x[-1]
     if make_comparable:
         x = make_comparable(x)
 

tests/logprob/test_transforms.py

@@ -991,14 +991,13 @@ def test_absolute_rv_transform(test_val):
         (pt.erfc, ErfcTransform()),
         (pt.erfcx, ErfcxTransform()),
         (pt.sinh, SinhTransform()),
-        (pt.cosh, CoshTransform()),
         (pt.tanh, TanhTransform()),
         (pt.arcsinh, ArcsinhTransform()),
         (pt.arccosh, ArccoshTransform()),
         (pt.arctanh, ArctanhTransform()),
     ],
 )
-def test_extra_rv_transforms(pt_transform, transform):
+def test_extra_bijective_rv_transforms(pt_transform, transform):
     base_rv = pt.random.normal(
         0.5, 1, name="base_rv"
     )  # Something not centered around 0 is usually better
@@ -1011,7 +1010,29 @@ def test_extra_rv_transforms(pt_transform, transform):
 
     vv_test = np.array(0.25)  # Arbitrary test value
     np.testing.assert_allclose(
-        rv_logp.eval({vv: vv_test}), np.nan_to_num(expected_logp.eval({vv: vv_test}), nan=-np.inf)
+        rv_logp.eval({vv: vv_test}),
+        np.nan_to_num(expected_logp.eval({vv: vv_test}), nan=-np.inf),
+    )
+
+
+def test_cosh_rv_transform():
+    # Something not centered around 0 is usually better
+    base_rv = pt.random.normal(0.5, 1, size=(2,), name="base_rv")
+    rv = pt.cosh(base_rv)
+
+    vv = rv.clone()
+    rv_logp = logp(rv, vv)
+
+    transform = CoshTransform()
+    [back_neg, back_pos] = transform.backward(vv)
+    expected_logp = pt.logaddexp(
+        logp(base_rv, back_neg), logp(base_rv, back_pos)
+    ) + transform.log_jac_det(vv)
+
+    vv_test = np.array([0.25, 1.5])
+    np.testing.assert_allclose(
+        rv_logp.eval({vv: vv_test}),
+        np.nan_to_num(expected_logp.eval({vv: vv_test}), nan=-np.inf),
     )