Allow do interventions to reference intervened variable

Author: ricardoV94

pymc_experimental/model_transform/conditioning.py

@@ -1,7 +1,9 @@
+import warnings
 from typing import Any, Dict, List, Sequence, Union
 
 from pymc import Model
 from pymc.pytensorf import _replace_vars_in_graphs
+from pytensor.graph import ancestors
 from pytensor.tensor import TensorVariable
 
 from pymc_experimental.model_transform.basic import prune_vars_detached_from_observed
@@ -188,7 +190,15 @@ def do(
         # Just a sanity check
         assert model_var in fgraph.variables
 
-        intervention.name = model_var.name
+        # If the intervention references the original variable we must give it a different name
+        if model_var in ancestors([intervention]):
+            intervention.name = f"do_{model_var.name}"
+            warnings.warn(
+                f"Intervention expression references the variable that is being intervened: {model_var.name}. "
+                f"Intervention will be given the name: {intervention.name}"
+            )
+        else:
+            intervention.name = model_var.name
         dims = extract_dims(model_var)
         # If there are any RVs in the graph we introduce the intervention as a deterministic
         if rvs_in_graph([intervention]):

pymc_experimental/tests/model_transform/test_conditioning.py

@@ -214,3 +214,16 @@ def test_do_prune(prune):
         assert set(do_m.named_vars) == {"x1", "z", "llike"}
     else:
         assert set(do_m.named_vars) == orig_named_vars
+
+
+def test_do_self_reference():
+    """Check we can replace a variable by an expression that refers to the same variable."""
+    with pm.Model() as m:
+        x = pm.Normal("x", 0, 1)
+
+    with pytest.warns(
+        UserWarning,
+        match="Intervention expression references the variable that is being intervened",
+    ):
+        new_x = do(m, {x: x + 100})["do_x"]
+    assert pm.draw(new_x) > 50