Update warning for transformed variables that depend on marginalized variable

Warning is now only issued for IntervalTransforms

Author: ricardoV94

pymc_experimental/marginal_model.py

@@ -5,8 +5,10 @@
 import pytensor.tensor as pt
 from pymc import SymbolicRandomVariable
 from pymc.distributions.discrete import Bernoulli, Categorical, DiscreteUniform
+from pymc.distributions.transforms import Chain
 from pymc.logprob.abstract import _get_measurable_outputs, _logprob
 from pymc.logprob.joint_logprob import factorized_joint_logprob
+from pymc.logprob.transforms import IntervalTransform
 from pymc.model import Model
 from pymc.pytensorf import constant_fold, inputvars
 from pytensor import Mode
@@ -171,11 +173,19 @@ def _marginalize(self, user_warnings=False):
                 if old_rv in self.basic_RVs:
                     self._transfer_rv_mappings(old_rv, new_rv)
                     if user_warnings:
+                        # Interval transforms for dependent variable won't work for non-constant bounds because
+                        # the RV inputs are now different and may depend on another RV that also depends on the
+                        # same marginalized RV
                         transform = self.rvs_to_transforms[new_rv]
-                        if transform is not None:
+                        if isinstance(transform, IntervalTransform) or (
+                            isinstance(transform, Chain)
+                            and any(
+                                isinstance(tr, IntervalTransform) for tr in transform.transform_list
+                            )
+                        ):
                             warnings.warn(
-                                "Transforms for variables that depend on marginalized RVs are currently not working, "
-                                f"rv={new_rv}, transform={transform}",
+                                f"The transform {transform} for the variable {old_rv}, which depends on the "
+                                f"marginalized {rv_to_marginalize} may no longer work if bounds depended on other variables.",
                                 UserWarning,
                             )
         return self

pymc_experimental/tests/test_marginal_model.py

@@ -346,16 +346,19 @@ def test_not_supported_marginalized_deterministic_and_potential():
     "transform, expected_warning",
     (
         (None, does_not_warn()),
-        pytest.param(
-            UNSET,
+        (UNSET, does_not_warn()),
+        (transforms.log, does_not_warn()),
+        (transforms.Chain([transforms.log, transforms.logodds]), does_not_warn()),
+        (
+            transforms.Interval(0, 1),
             pytest.warns(
-                UserWarning, match="Transforms for variables that depend on marginalized RVs"
+                UserWarning, match="which depends on the marginalized idx may no longer work"
             ),
         ),
-        pytest.param(
-            transforms.log,
+        (
+            transforms.Chain([transforms.log, transforms.Interval(0, 1)]),
             pytest.warns(
-                UserWarning, match="Transforms for variables that depend on marginalized RVs"
+                UserWarning, match="which depends on the marginalized idx may no longer work"
             ),
         ),
     ),
@@ -398,5 +401,9 @@ def test_marginalized_transforms(transform, expected_warning):
 
     ip = m.initial_point()
     if transform is not None:
-        assert "sigma_log__" in ip
+        if transform is UNSET:
+            transform_name = "log"
+        else:
+            transform_name = transform.name
+        assert f"sigma_{transform_name}__" in ip
     np.testing.assert_allclose(m.compile_logp()(ip), m_ref.compile_logp()(ip))