Fix bug in model_fgraph where RNGs weren't being copied as advertised

Author: ricardoV94

pymc_experimental/tests/utils/test_model_fgraph.py

@@ -118,7 +118,7 @@ def test_data(inline_views):
     assert m_new.rvs_to_values[m_new["obs"]] is m_new["y"]
 
     # Shared rng shared variables are not preserved
-    m_new["b1"].owner.inputs[0].container is not m_old["b1"].owner.inputs[0].container
+    assert m_new["b1"].owner.inputs[0].container is not m_old["b1"].owner.inputs[0].container
 
     with m_old:
         pm.set_data({"x": [100.0, 200.0]}, coords={"test_dim": range(2)})

pymc_experimental/utils/model_fgraph.py

@@ -4,7 +4,7 @@
 from pymc.logprob.transforms import RVTransform
 from pymc.model import Model
 from pymc.pytensorf import find_rng_nodes
-from pytensor import Variable
+from pytensor import Variable, shared
 from pytensor.graph import Apply, FunctionGraph, Op, node_rewriter
 from pytensor.graph.rewriting.basic import out2in
 from pytensor.scalar import Identity
@@ -184,8 +184,7 @@ def fgraph_from_model(
 
     # Replace RNG nodes so that seeding does not interfere with old model
     for rng in find_rng_nodes(model_vars):
-        new_rng = rng.clone()
-        new_rng.set_value(rng.get_value(borrow=False))
+        new_rng = shared(rng.get_value(borrow=False))
         memo[rng] = new_rng
 
     fgraph = FunctionGraph(