Calling rvs_to_value_vars modifies model.free_RVs

[michaelosthege]

This error was found while working on #5170:

Please provide a minimal, self-contained, and reproducible example.

with pm.Model() as pmodel:
    hyper = pm.LogNormal("hyper", mu=0)
    group = pm.LogNormal("group", mu=at.log(hyper))

    # We add potentials or deterministics that are not in topological order
    pm.Potential('group_pot', group)
    pm.Potential("hyper_pot", hyper)

    before = len(FunctionGraph(outputs=pmodel.free_RVs).toposort())

    # This call will change the model free_RVs in place!
    # rvs_to_value_vars(pmodel.deterministics, apply_transforms=True)
    rvs_to_value_vars(pmodel.potentials, apply_transforms=True)

    after = len(FunctionGraph(outputs=pmodel.free_RVs).toposort())

    assert before == after

Traceback (most recent call last):
  File "/home/ricardo/Documents/Projects/pymc3-venv/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3418, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-9-66cc9d4cc5c8>", line 37, in <module>
    assert before == after
AssertionError

The same happens with the Deterministics, and can be checked by flipping the commented lines and commenting those referring to Potentials

Versions and main components

• PyMC/PyMC3 Version: main
• Aesara/Theano Version: 2.2.6

[ricardoV94]

It seems that with some combinations of transforms and special ordering of variables the function rvs_to_value_vars changes the model.free_RVs in place.

CC @brandonwillard

[brandonwillard]

The problem is that the model variables named "hyper" and "group" (i.e. the ones in model.free_RVs) are both transformed variables, they're nested, and they both end up in the replacements dict due to transform_replacements (because they need to be replaced with their transformed counterparts, of course).

When replace_rvs_in_graphs is called later, the initial_replacements argument consists of the replacements gathered by transform_replacements, and those model variables are not cloned by clone_get_equiv because they're already in the equiv dict.

The end result is that fg contains the original "group" variable, and, when fg.replace_all is called, it works in-place on it—replacing the subgraph containing "hyper" with its transformed value variable.

Cloning the graphs before calling rvs_to_value_vars should fix that, e.g.:

from aesara.graph.basic import Constant, clone_get_equiv, graph_inputs

inputs = [i for i in graph_inputs(pmodel.potentials) if not isinstance(i, Constant)]
equiv = clone_get_equiv(inputs, pmodel.potentials, False, False, {})

rv_val_res = rvs_to_value_vars([equiv[n] for n in pmodel.potentials], apply_transforms=True)