clone_replace's share_inputs doesn't consider RandomVariables as inputs

[lucianopaz]

Description of your problem or feature request

I have a biggish hierarchical model and I want to test an intervention where I force some of the random variables to be exactly equal to 0. There are many ways in which I can do this, but I tried to use clone_replace to replace the intervened random variables with zeros of the correct shape and dtype. I chose not to do this by setting givens when I compiled the function, because I actually want my function to output both conditions: with and without interventions.

The problem I faced was that clone_replace made copies of the random variables that weren't being changed (the random generators and other constant inputs were not cloned because those were the graph inputs), and then when I compiled a function to draw samples, the intervention model resampled the cloned variables from their prior. Here's a minimal example:

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

a = at.random.normal(loc=3, scale=0.01, name="a", size=2)
b = at.random.normal(loc=1, scale=0.01, name="b", size=(2, 2))
c = at.random.normal(loc=100, scale=0.01, name="c", size=(2, 2, 2))
d = pm.Normal.dist(mu=(a + b + c).flatten(), sigma=0.01)
d.name = "d"
d_clone = aesara.graph.basic.clone_replace(
    [d], replace={c: at.zeros(c.shape, dtype=c.dtype)}, share_inputs=True
)
f = aesara.function([a, b], d_clone, on_unused_input="ignore")
f(a=np.zeros(2), b=np.zeros((2, 2)))

Which prints:

[array([3.98356155, 3.9887164 , 3.99668762, 4.01093951, 3.98949882, 0129146 , 3.9909093 , 4.04306206])]

instead of being centered around zero.

If instead of using clone_replace I use clone_get_equiv, I can put together something that works like I want it to:

a = at.random.normal(loc=3, scale=0.01, name="a", size=2)
b = at.random.normal(loc=1, scale=0.01, name="b", size=(2, 2))
c = at.random.normal(loc=100, scale=0.01, name="c", size=(2, 2, 2))
d = pm.Normal.dist(mu=(a + b + c).flatten(), sigma=0.01)
d.name = "d"
clone_map = aesara.graph.basic.clone_get_equiv([], [d],)
fg = aesara.graph.fg.FunctionGraph(None, [clone_map[d]], clone=False)
fg.replace_all(
    [
        (clone_map[a], a),
        (clone_map[b], b),
        (clone_map[c], at.zeros(clone_map[c].shape, dtype=clone_map[c].dtype)),
    ],
    import_missing=True,
)
aesara.function([a, b], clone_map[d])(a=np.zeros(2), b=np.zeros((2, 2)))

Which prints out

array([ 0.00141859, -0.02443352, -0.00727199,  0.00841104, -0.00885378, -0.00093795,  0.00688666, -0.00521271])

I understand now why variables a and b are not shared after calling clone_replace but I think that it would be useful to either:

• Have clone_replace return the mapping between original and cloned nodes
• Be able to pass an inputs list to clone_replace that indicates which nodes should be considered inputs to be shared.

Versions and main components

• Aesara version: 2.6.6
• Aesara config (python -c "import aesara; print(aesara.config)")
• Python version: 3.9
• Operating system: Ubuntu 18.04
• How did you install Aesara: (conda/pip) conda

[rlouf]

Could you print the graph in both cases using aesara.dprint?

[lucianopaz]

I'll print d_clone and clone_map[d].

Using clone_replace:

normal_rv{0, (0, 0), floatX, False}.1 [id A] 'd'
 |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD74111F740>) [id B]
 |TensorConstant{[]} [id C]
 |TensorConstant{11} [id D]
 |Reshape{1} [id E]
 | |Elemwise{add,no_inplace} [id F]
 | | |InplaceDimShuffle{x,0,1} [id G]
 | | | |Elemwise{add,no_inplace} [id H]
 | | |   |InplaceDimShuffle{x,0} [id I]
 | | |   | |normal_rv{0, (0, 0), floatX, False}.1 [id J] 'a'
 | | |   |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD74287C200>) [id K]
 | | |   |   |TensorConstant{(1,) of 2} [id L]
 | | |   |   |TensorConstant{11} [id M]
 | | |   |   |TensorConstant{3} [id N]
 | | |   |   |TensorConstant{0.01} [id O]
 | | |   |normal_rv{0, (0, 0), floatX, False}.1 [id P] 'b'
 | | |     |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD74287CAC0>) [id Q]
 | | |     |TensorConstant{(2,) of 2} [id R]
 | | |     |TensorConstant{11} [id S]
 | | |     |TensorConstant{1} [id T]
 | | |     |TensorConstant{0.01} [id U]
 | | |Alloc [id V]
 | |   |TensorConstant{0.0} [id W]
 | |   |Subtensor{int64} [id X]
 | |   | |Shape [id Y]
 | |   | | |normal_rv{0, (0, 0), floatX, False}.1 [id Z] 'c'
 | |   | |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD741104740>) [id BA]
 | |   | |   |TensorConstant{(3,) of 2} [id BB]
 | |   | |   |TensorConstant{11} [id BC]
 | |   | |   |TensorConstant{100} [id BD]
 | |   | |   |TensorConstant{0.01} [id BE]
 | |   | |ScalarConstant{0} [id BF]
 | |   |Subtensor{int64} [id BG]
 | |   | |Shape [id Y]
 | |   | |ScalarConstant{1} [id BH]
 | |   |Subtensor{int64} [id BI]
 | |     |Shape [id Y]
 | |     |ScalarConstant{2} [id BJ]
 | |TensorConstant{(1,) of -1} [id BK]
 |TensorConstant{0.01} [id BL]

Using clone_get_equiv:

normal_rv{0, (0, 0), floatX, False}.1 [id A] 'd'
 |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD74350E040>) [id B]
 |TensorConstant{[]} [id C]
 |TensorConstant{11} [id D]
 |Reshape{1} [id E]
 | |Elemwise{add,no_inplace} [id F]
 | | |InplaceDimShuffle{x,0,1} [id G]
 | | | |Elemwise{add,no_inplace} [id H]
 | | |   |InplaceDimShuffle{x,0} [id I]
 | | |   | |normal_rv{0, (0, 0), floatX, False}.1 [id J] 'a'
 | | |   |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD74111FE40>) [id K]
 | | |   |   |TensorConstant{(1,) of 2} [id L]
 | | |   |   |TensorConstant{11} [id M]
 | | |   |   |TensorConstant{3} [id N]
 | | |   |   |TensorConstant{0.01} [id O]
 | | |   |normal_rv{0, (0, 0), floatX, False}.1 [id P] 'b'
 | | |     |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD741104820>) [id Q]
 | | |     |TensorConstant{(2,) of 2} [id R]
 | | |     |TensorConstant{11} [id S]
 | | |     |TensorConstant{1} [id T]
 | | |     |TensorConstant{0.01} [id U]
 | | |Alloc [id V]
 | |   |TensorConstant{0.0} [id W]
 | |   |Subtensor{int64} [id X]
 | |   | |Shape [id Y]
 | |   | | |normal_rv{0, (0, 0), floatX, False}.1 [id Z] 'c'
 | |   | |   |RandomGeneratorSharedVariable(<Generator(PCG64) at 0x7FD74287CBA0>) [id BA]
 | |   | |   |TensorConstant{(3,) of 2} [id BB]
 | |   | |   |TensorConstant{11} [id BC]
 | |   | |   |TensorConstant{100} [id BD]
 | |   | |   |TensorConstant{0.01} [id BE]
 | |   | |ScalarConstant{0} [id BF]
 | |   |Subtensor{int64} [id BG]
 | |   | |Shape [id Y]
 | |   | |ScalarConstant{1} [id BH]
 | |   |Subtensor{int64} [id BI]
 | |     |Shape [id Y]
 | |     |ScalarConstant{2} [id BJ]
 | |TensorConstant{(1,) of -1} [id BK]
 |TensorConstant{0.01} [id BL]

[lucianopaz]

Furthermore, if I try to compile the function with clone_replace and don't use on_missing_inputs="ignore", I get an error:

a = at.random.normal(loc=3, scale=0.01, name="a", size=2)
b = at.random.normal(loc=1, scale=0.01, name="b", size=(2, 2))
c = at.random.normal(loc=100, scale=0.01, name="c", size=(2, 2, 2))
d = pm.Normal.dist(mu=(a + b + c).flatten(), sigma=0.01)
d.name = "d"
d_clone = aesara.graph.basic.clone_replace(
    [d], replace={c: at.zeros(c.shape, dtype=c.dtype)}, share_inputs=True
)
f = aesara.function([a, b], d_clone)

Traceback:

UnusedInputError                          Traceback (most recent call last)
/tmp/ipykernel_21661/2516318076.py in <module>
      7     [d], replace={c: at.zeros(c.shape, dtype=c.dtype)}, share_inputs=True
      8 )
----> 9 f = aesara.function([a, b], d_clone)

~/anaconda3/lib/python3.9/site-packages/aesara/compile/function/__init__.py in function(inputs, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input)
    315         # note: pfunc will also call orig_function -- orig_function is
    316         #      a choke point that all compilation must pass through
--> 317         fn = pfunc(
    318             params=inputs,
    319             outputs=outputs,

~/anaconda3/lib/python3.9/site-packages/aesara/compile/function/pfunc.py in pfunc(params, outputs, mode, updates, givens, no_default_updates, accept_inplace, name, rebuild_strict, allow_input_downcast, profile, on_unused_input, output_keys)
    361     )
    362 
--> 363     return orig_function(
    364         inputs,
    365         cloned_outputs,

~/anaconda3/lib/python3.9/site-packages/aesara/compile/function/types.py in orig_function(inputs, outputs, mode, accept_inplace, name, profile, on_unused_input, output_keys)
   1723     try:
   1724         Maker = getattr(mode, "function_maker", FunctionMaker)
-> 1725         m = Maker(
   1726             inputs,
   1727             outputs,

~/anaconda3/lib/python3.9/site-packages/aesara/compile/function/types.py in __init__(self, inputs, outputs, mode, accept_inplace, function_builder, profile, on_unused_input, fgraph, output_keys, name)
   1430 
   1431         # Check if some input variables are unused
-> 1432         self.check_unused_inputs(inputs, outputs, on_unused_input)
   1433 
   1434         indices = [[input, None, [input]] for input in inputs]

~/anaconda3/lib/python3.9/site-packages/aesara/compile/function/types.py in check_unused_inputs(inputs, outputs, on_unused_input)
   1371                     )
   1372                 elif on_unused_input == "raise":
-> 1373                     raise UnusedInputError(msg % (inputs.index(i), i.variable, err_msg))
   1374                 else:
   1375                     raise ValueError(

UnusedInputError: aesara.function was asked to create a function computing outputs given certain inputs, but the provided input variable at index 0 is not part of the computational graph needed to compute the outputs: a.
To make this error into a warning, you can pass the parameter on_unused_input='warn' to aesara.function. To disable it completely, use on_unused_input='ignore'.

[brandonwillard]

I understand now why variables a and b are not shared after calling clone_replace but I think that it would be useful to either:

• Have clone_replace return the mapping between original and cloned nodes
• Be able to pass an inputs list to clone_replace that indicates which nodes should be considered inputs to be shared.

clone_replace is a somewhat restrictive form of cloning and replacement in a graph. As you noticed, clone_get_equiv is a more robust means of performing replacements.

The one thing clone_replace doesn't really do is preserve the identities/equivalences of subgraphs. During compilation, those identity relationships are reconstructed by the "merge" rewrite, which will replace the output variables of equivalent Apply nodes (i.e. nodes using equivalent Ops and inputs) with a single "canonical" instance of those outputs (i.e. choose one Apply node to represent the equivalence class and only ever use its outputs).

I've been going over all this recently and planning some core graph changes that will preserve equivalence and identities at a much lower level (e.g. cache Apply instances), potentially removing the need for a distinct "merging" rewrite altogether. These changes would help facilitate things like #1082 and make replacements like yours much more straightforward and consistent.

[lucianopaz]

I understand that the merge rewrite will replace the output variables or equivalent Apply nodes. But my problem is that when I want to compile the function by specifying inputs using the original variables, I can't. I think that it would be nice to be able to tell clone_replace which nodes it should consider as inputs, at least when it says that the inputs will be shared across the original and cloned versions.

[ricardoV94]

We probably can do this already with FunctionGraph(inputs, outputs, clone=True) and then replacing it in place. I agree it would be nice to have this functionality in the more user friendly clone_replace if there is not good reason not to.

[brandonwillard]

But my problem is that when I want to compile the function by specifying inputs using the original variables, I can't. I think that it would be nice to be able to tell clone_replace which nodes it should consider as inputs, at least when it says that the inputs will be shared across the original and cloned versions.

There are two fundamental issues here:

1. you're calling a graph "cloning" function, but you don't actually want to clone the entire graph, and
2. your "inputs" are not inputs according to Aesara; they're sub-graphs.

Both of these issues imply that your first problem (i.e. not being able to using the original variables) is inherent to the approach, because, unless your "inputs" are actual inputs according to Aesara, the input-copying options (e.g. copy_inputs_over) aren't relevant. If they aren't inputs, and are instead sub-graphs, then they'll be cloned—by definition—and no longer the original variables. (The "merging" I mentioned earlier can settle this during compilation, though.)

As you've noticed, clone_get_equiv allows one to get much closer to the functionality you want. The naming and docstrings of these two cloning-related functions aren't great, so those could use some work, but it wouldn't help to make functional changes to clone_replace—especially not if they make it work more like clone_get_equiv, or vice versa. Remember, clone_replace is just an alternative interface to rebuild_collect_shared, and rebuild_collect_shared performs some very specific operations that are needed in Aesara.

Also, if you use the memo argument to clone_get_equiv, you can specify which terms are not cloned (e.g. map an existing Variable to itself and it will remain in the graph). Likewise, clone_get_equiv has an inputs argument that can be used to restrict which terms are cloned and/or considered as inputs.