Support "transitive" Scan log-probabilities

[brandonwillard]

Currently, Scan log-probability support only handles cases in which the MeasurableVariable is created inside the body/step function of the Scan, and not when the body/step function simply references a MeasurableVariable that is being iterated over by the Scan.

For example, the following is not supported:

import aesara
import aesara.tensor as at

from aeppl.joint_logprob import factorized_joint_logprob


srng = at.random.RandomStream(seed=2320)
N = 10

Y_rv = srng.normal(0, 1, size=N, name="Y")


def step_fn(y_t):
    return y_t


Y_1T_rv, _ = aesara.scan(
    fn=step_fn,
    sequences=[Y_rv],
    strict=True,
)

y_vv = Y_1T_rv.clone()
y_vv.name = "y"

logp_parts = factorized_joint_logprob({Y_1T_rv: y_vv})

This example is very trivial, but, if we change step_fn so that it performs a supported, measurable operation on y_t (e.g. indexing a mixture), it wouldn't work for the same reason.

When a value is assigned to Scan output terms like Y_1T_rv, we could "push" the relevant sequences inputs into the step function. In other words, we could construct the type of graph we currently handle.

Working from the example above, we would rewrite the Scan into something like the following:

# Apply a rewrite like `local_rv_size_lift` to get properly `size`-broadcasted parameters
# in a new variable `new_Y_rv`
mu_bcast, sigma_bcast = new_Y_rv.owner.inputs[3:]

def new_step_fn(mu_t, sigma_t)
    return Y_rv.owner.op(mu_t, sigma_t, name="Y[t]")

new_Y_1T_rv, _ = aesara.scan(
    fn=new_step_fn,
    sequences=[mu_bcast, sigma_bcast],
    strict=True,
)