Use callable for default_update

[ricardoV94]

The default_update is set by RandomStream so that users don't have to manually specify update expressions for RandomState/Generators in Aesara functions.

However, this can be problematic when manipulating graphs, because they create a hidden dependency that is not revealed until the final function is actually compiled.

import aesara
import aesara.tensor as at

size = at.constant(2)
srng = at.random.RandomStream(0)
x = srng.normal(size=size)
(x,) = aesara.graph.clone_replace([x], {size: at.constant(4)})

f = aesara.function([], x)
print(f()[2:4], f()[:2])  # [0.73595567 0.00587704] [0.73595567 0.00587704]

In that example the shared variable used by the cloned x has the wrong update expression, causing the bit generator to only advance 2 steps every call, instead of 4. It also produces a wasteful graph that takes 2 unused random draws in the background just to update the shared variable. There are other caveats when manipulating graphs with RandomVariables, but this one seems unnecessary.

One option would be to replace the stateful default_update by a callable that returns the correct update expression for a given node. This is the approach taken in PyMC for example (https://github.com/pymc-devs/pymc/blob/4c92adf9720f6578e51b9ef21e33c29871c67a83/pymc/aesaraf.py#L928-L933), and it is also similar to the approach taken by Aeppl transforms to avoid new "surprising dependencies" at runtime (https://github.com/aesara-devs/aeppl/blob/b4304fa60979625450b79fe4cfb88426bc3d73e8/aeppl/transforms.py#L426-L429).

Instead of setting x.owner.inputs[0].default_update = x.owner.outputs[0], RandomStream would set x.owner.default_update = lambda node: (node.inputs[0]: node.outputs[0]). The internal machinery would need to be updated to look for default_updates in nodes instead of SharedVariables, but otherwise I don't think it would be too much trouble.

An altogether different solution was attempted in #896

[brandonwillard]

One option would be to replace the stateful default_update by a callable that returns the correct update expression for a given node. This is the approach taken in PyMC for example (https://github.com/pymc-devs/pymc/blob/4c92adf9720f6578e51b9ef21e33c29871c67a83/pymc/aesaraf.py#L928-L933), and it is also similar to the approach taken by Aeppl transforms to avoid new "surprising dependencies" at runtime (https://github.com/aesara-devs/aeppl/blob/b4304fa60979625450b79fe4cfb88426bc3d73e8/aeppl/transforms.py#L426-L429).

It seems like you're talking about making default_update hold a graph constructor that would assemble the update graphs given some undetermined inputs. If you're talking about that, then there would be a lot more involved and it would probably be easier to fix clone_replace and/or whatever else is preventing default_update graphs from being correctly cloned.

Basically, anything that would (re)construct a graph can be loosely called "cloning", which means the callable would only encode the construction of the graph; however, I don't think the challenge lies there, since "proper" cloning handles all this and more just fine—when applied correctly. Likewise, this callable approach would need to determine the callable's inputs and orchestrate the use of the callables, and that's where one would likely end up effectively reproducing all the same logic as proper cloning.

I also don't understand the connection with those links and the use of a callable and/or default_updates.

Instead of setting x.owner.inputs[0].default_update = x.owner.outputs[0], RandomStream would set x.owner.default_update = lambda node: (node.inputs[0]: node.outputs[0]). The internal machinery would need to be updated to look for default_updates in nodes instead of SharedVariables, but otherwise I don't think it would be too much trouble.

This only applies to one type of update graph, so it's not a viable solution to the general default_update + cloning problem. It would be very convenient for its subset of cases, though.

An altogether different solution was attempted in #896

Note that #896 was not just about fixing the default_update + cloning issue. Also, I closed that PR because the actual issue(s) we were trying to solve (e.g. #579) were not properly addressed by #896 alone. Regardless, I think the default_update + cloning issue could be addressed directly without too much effort.

[ricardoV94]

This only applies to one type of update graph

I know default update is more general, but do we really need it to be? It was clearly added specifically for the case of RandomVariables where it's enough to map one of the node inputs to its own output.

Perhaps it's fine to restrict the default update to those cases and leave the manual updates for anything less local.

[brandonwillard]

Perhaps it's fine to restrict the default update to those cases and leave the manual updates for anything less local.

We can't do that without removing updates entirely.

[ricardoV94]

Why does it affect updates other than those specified via default_update?