Fix bug in nested replacement of value vars

Author: ricardoV94

pymc/aesaraf.py

@@ -236,6 +236,7 @@ def expand_replace(var):
             new_nodes.extend(replacement_fn(var, replacements))
         return new_nodes
 
+    # This iteration populates the replacements
     for var in walk_model(graphs, expand_fn=expand_replace, **kwargs):
         pass
 
@@ -250,7 +251,15 @@ def expand_replace(var):
             clone=False,
         )
 
-        fg.replace_all(replacements.items(), import_missing=True)
+        # replacements have to be done in reverse topological order so that nested
+        # expressions get recursively replaced correctly
+        toposort = fg.toposort()
+        sorted_replacements = sorted(
+            tuple(replacements.items()),
+            key=lambda pair: toposort.index(pair[0].owner),
+            reverse=True,
+        )
+        fg.replace_all(sorted_replacements, import_missing=True)
 
         graphs = list(fg.outputs)
 

pymc/tests/distributions/test_logprob.py

@@ -52,12 +52,7 @@
     logp,
 )
 from pymc.model import Model, Potential
-from pymc.tests.helpers import select_by_precision
-
-
-def assert_no_rvs(var):
-    assert not any(isinstance(v.owner.op, RandomVariable) for v in ancestors([var]) if v.owner)
-    return var
+from pymc.tests.helpers import assert_no_rvs, select_by_precision
 
 
 def test_get_scaling():

pymc/tests/helpers.py

@@ -24,8 +24,10 @@
 import numpy.random as nr
 
 from aesara.gradient import verify_grad as at_verify_grad
+from aesara.graph import ancestors
 from aesara.graph.rewriting.basic import in2out
 from aesara.sandbox.rng_mrg import MRG_RandomStream as RandomStream
+from aesara.tensor.random.op import RandomVariable
 
 import pymc as pm
 
@@ -218,3 +220,8 @@ def continuous_steps(self, step, step_kwargs):
             assert {m.rvs_to_values[c1], m.rvs_to_values[c2]} == set(
                 step([c1, c2], **step_kwargs).vars
             )
+
+
+def assert_no_rvs(var):
+    assert not any(isinstance(v.owner.op, RandomVariable) for v in ancestors([var]) if v.owner)
+    return var

pymc/tests/test_aesaraf.py

@@ -47,6 +47,7 @@
 from pymc.distributions.distribution import SymbolicRandomVariable
 from pymc.distributions.transforms import Interval
 from pymc.exceptions import NotConstantValueError
+from pymc.tests.helpers import assert_no_rvs
 from pymc.vartypes import int_types
 
 
@@ -632,3 +633,56 @@ def test_no_change_inplace(self):
             after = aesara.clone_replace(m.free_RVs)
 
             assert equal_computations(before, after)
+
+    @pytest.mark.parametrize("reversed", (False, True))
+    def test_interdependent_transformed_rvs(self, reversed):
+        # Test that nested transformed variables, whose transformed values depend on other
+        # RVs are properly replaced
+        with pm.Model() as m:
+            transform = pm.distributions.transforms.Interval(
+                bounds_fn=lambda *inputs: (inputs[-2], inputs[-1])
+            )
+            x = pm.Uniform("x", lower=0, upper=1, transform=transform)
+            y = pm.Uniform("y", lower=0, upper=x, transform=transform)
+            z = pm.Uniform("z", lower=0, upper=y, transform=transform)
+            w = pm.Uniform("w", lower=0, upper=z, transform=transform)
+
+        rvs = [x, y, z, w]
+        if reversed:
+            rvs = rvs[::-1]
+
+        transform_values = rvs_to_value_vars(rvs)
+
+        for transform_value in transform_values:
+            assert_no_rvs(transform_value)
+
+        if reversed:
+            transform_values = transform_values[::-1]
+        transform_values_fn = m.compile_fn(transform_values, point_fn=False)
+
+        x_interval_test_value = np.random.rand()
+        y_interval_test_value = np.random.rand()
+        z_interval_test_value = np.random.rand()
+        w_interval_test_value = np.random.rand()
+
+        # The 3 Nones correspond to unused rng, dtype and size arguments
+        expected_x = transform.backward(x_interval_test_value, None, None, None, 0, 1).eval()
+        expected_y = transform.backward(
+            y_interval_test_value, None, None, None, 0, expected_x
+        ).eval()
+        expected_z = transform.backward(
+            z_interval_test_value, None, None, None, 0, expected_y
+        ).eval()
+        expected_w = transform.backward(
+            w_interval_test_value, None, None, None, 0, expected_z
+        ).eval()
+
+        np.testing.assert_allclose(
+            transform_values_fn(
+                x_interval__=x_interval_test_value,
+                y_interval__=y_interval_test_value,
+                z_interval__=z_interval_test_value,
+                w_interval__=w_interval_test_value,
+            ),
+            [expected_x, expected_y, expected_z, expected_w],
+        )