Fix some usability issues surrounding shared RNG objects

[brandonwillard]

Work on pymc-devs/pymc#4729 has highlighted a potential confusion involving RandomStateSharedVariables and in-place optimizations like random_make_inplace.

The problem is that in-place optimizations won't be performed on RandomStateSharedVariables that do not have updates specified, since they're protected by the FunctionGraph Supervisor feature set up by aesara.compile.function.types.std_fgraph.

In order to get around this, one can add/use the default_update property on the RandomStateSharedVariable (e.g. see here and RandomStream.gen); however, this isn't set automatically, so it can lead to confusion.

Let's consider adding the default_update property automatically. This could probably be done in RandomVariable.make_node, but we first need to consider whether or not this will have other repercussions/restrictions/etc.

[ricardoV94]

whether or not this will have other repercussions/restrictions/etc.

One possible limitation concerns variable replacements, an example of which emerged here: pymc-devs/pymc#4903 (comment)

Maybe we could solve it at compile time instead? Is it possible for a rewrite to manipulate the default_update of Shared variables?

@local_optimizer([RandomVariable])
def rv_default_update(node, fgraph):    
    rng = node.owner.inputs[0]
    if not hasattr(rng, "default_update"):
        rng.default_update = rv.owner.outputs[0]

[brandonwillard]

Maybe we could solve it at compile time instead? Is it possible for a rewrite to manipulate the default_update of Shared variables?

That's a great question, and I'm confident that this idea will be important at some point.

One of the issues with this is that the updates are returned to the user as an object that's independent from the graph that uses the updates, so there's always the chance that users will manually specify updates when calling aesara.function. I'm assuming we can ignore such updates when a graph has been rewritten to—say—remove the relevant nodes from the compiled graph.

Aside from that, I believe that updates-generating Ops like Scan can be rewritten to remove the terms that are directly associated with the corresponding updates. For instance, Scans can be refactored to remove the outer-inputs and inner graphs elements that correspond to updates; we've come across this in aesara-devs/aeppl#24.

[brandonwillard]

Following up on a Gitter conversation about this, it's really the interaction between Scan and a special .update attribute that makes RandomVariables work as expected within a Scan.

Those attributes are added to the outputs of RandomVariables by RandomStream, which is the main reason why the RandomVariables produced by RandomStream work with Scan. Scan picks up the values in that attribute and adds them to the updates it produces.

We need to do something so that effectively arbitrary use of shared RNG objects is possible. For instance, we could change Scan so that it performs the same operations that it currently does when .updates is present, but without the need for a .updates.

That and/or we should make it very clear that RandomStream is the interface for generating RandomVariable outputs.