Support RandomState Variable types in Scan

[brandonwillard]

Currently, it's not possible to manually manage the RNG state evolution of RandomVariables in Scan, and it's due to limited Variable type support.

Fixing this would take us a large step toward removing the need to use shared variables and updates when making combined use of Scans and RandomVariables.

For example,

import aesara
import aesara.tensor as at

from aesara.tensor.random.type import random_generator_type


def inner_fn(prev_rng):
    res = at.random.normal(rng=prev_rng)
    new_rng = res.owner.outputs[0]
    return new_rng


initial_rng = random_generator_type()

out, _ = aesara.scan(
    inner_fn, outputs_info={"initial": initial_rng, "taps": [-1]}, n_steps=10
)
# TypeError: Tensor type field must be a TensorType; found <class 'aesara.tensor.random.type.RandomGeneratorType'>.

Since it won't be easy to return a collection of all intermediate RNG state variables, we could instead return just the last one (if any). These are the kinds of things that need to be determined in order to implement this.

It's also possible that we could add special handling that reduces the boilerplate. For example, we could make it possible to omit the explicit RNG return statements, so that the following is sufficient to produce distinct samples on each iteration:

def inner_fn(prev_rng):
    res = at.random.normal(rng=prev_rng)
    return res


initial_rng = random_generator_type()

out, _ = aesara.scan(
    inner_fn, non_sequences=[initial_rng], n_steps=10
)

[ricardoV94]

It's also possible that we could add special handling that reduces the boilerplate. For example, we could make it possible to omit the explicit RNG return statements, so that the following is sufficient to produce distinct samples on each iteration:

How would that work behind the scenes?

I don't have a big issue with returning next rngs explicitly. We could also add a helper property to RandomVariables like .next_rng that returns the first output to make it more readable.

[brandonwillard]

It's also possible that we could add special handling that reduces the boilerplate. For example, we could make it possible to omit the explicit RNG return statements, so that the following is sufficient to produce distinct samples on each iteration:

How would that work behind the scenes?

We would construct a Scan that's equivalent to the first example.

[ricardoV94]

So first step here is to allow Scan to accept RandomGeneratorType variables as inputs?

[brandonwillard]

So first step here is to allow Scan to accept RandomGeneratorType variables as inputs?

Its base class, yes.