Repository: aesara-devs/aeppl (language: python)

Issue:
Implement a general mixture log-probability rewrite
This PR implements log-probabilities for general mixtures (e.g. heterogeneous `RandomVariable`s).

This implementation does not produce a **marginalized** log-probability.

It covers cases like the following:
```python
import aesara.tensor as at

from aeppl import joint_logprob


X_rv = at.random.normal(0, 1, name="X")
Y_rv = at.random.gamma(0.5, 0.5, name="Y")

p = at.scalar("p")
I_rv = at.random.bernoulli(p, name="I")

M_rv = at.stack([X_rv, Y_rv])[I_rv]

m = M_rv.type()
M_logp = joint_logprob(M_rv, {M_rv: m})
```

Instead of producing a log-probability graph that represents the computation `np.stack([logprob(X, obs), logprob(Y, obs)])[idx]`, which requires complete evaluation of `logprob(X, obs)` and `logprob(Y, obs)` even though `idx` will only select _some_ of those `logprob` values from each mixture component, this implementation will compute something like the following:
```python
res[idx == 0] = logprob(X[idx == 0], obs[idx == 0])
res[idx == 1] = logprob(Y[idx == 1], obs[idx == 1])
```

~Currently, it does all the necessary setup and parsing needed to construct the desired log-probability graphs; however, the actual construction steps aren't necessarily finished.  I still need to check/confirm a few things and&mdash;obviously&mdash;write tests.~


Discussion / hints:
The multivariate case is blocked by https://github.com/aesara-devs/aesara/pull/507 and https://github.com/aesara-devs/aesara/pull/509.

The latter issue is only a blocker if we want to prevent the evaluation of "unselected" mixture components.  We can always produce an array of stacked `logprob` values evaluated conditional on each mixture component and _then_ index that array with the mixing distribution; however, this is a prohibitively slow approach for large enough mixture models (e.g. almost anything involving time-varying mixture components used with non-trivially-sized datasets).
The two tests in this PR pass locally under the changes in https://github.com/aesara-devs/aesara/pull/516.
# [Codecov](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#19](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (a74b7ec) into [main](https://codecov.io/gh/aesara-devs/aeppl/commit/b07159bf69fa58e6923a7cac568356a80d1c1197?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (b07159b) will **increase** coverage by `2.27%`.
> The diff coverage is `93.75%`.

[![Impacted file tree graph](https://codecov.io/gh/aesara-devs/aeppl/pull/19/graphs/tree.svg?width=650&height=150&src=pr&token=L2i59LsFc0&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)

```diff
@@            Coverage Diff             @@
##             main      #19      +/-   ##
==========================================
+ Coverage   90.52%   92.80%   +2.27%     
==========================================
  Files           5        6       +1     
  Lines         802      917     +115     
  Branches      103      119      +16     
==========================================
+ Hits          726      851     +125     
+ Misses         44       30      -14     
- Partials       32       36       +4     
```


| [Impacted Files](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [aeppl/joint\_logprob.py](https://codecov.io/gh/aesara-devs/aeppl/pull/19/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-YWVwcGwvam9pbnRfbG9ncHJvYi5weQ==) | `98.00% <83.33%> (+6.33%)` | :arrow_up: |
| [aeppl/mixture.py](https://codecov.io/gh/aesara-devs/aeppl/pull/19/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-YWVwcGwvbWl4dHVyZS5weQ==) | `93.63% <93.63%> (ø)` | |
| [aeppl/opt.py](https://codecov.io/gh/aesara-devs/aeppl/pull/19/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-YWVwcGwvb3B0LnB5) | `82.60% <100.00%> (+22.00%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None). Last update [b07159b...a74b7ec](https://codecov.io/gh/aesara-devs/aeppl/pull/19?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None).

Is this working with categorical distributions as well? If so, we could add a test
> Is this working with categorical distributions as well? If so, we could add a test

Yes, definitely.

Fix patch:
diff --git a/aeppl/__init__.py b/aeppl/__init__.py
index 74d0a1c4..58c59c64 100644
--- a/aeppl/__init__.py
+++ b/aeppl/__init__.py
@@ -12,5 +12,6 @@
 # isort: off
 # Add optimizations to the DBs
 import aeppl.mixture
+import aeppl.scan
 
 # isort: on
diff --git a/aeppl/scan.py b/aeppl/scan.py
new file mode 100644
index 00000000..1a9a5e06
--- /dev/null
+++ b/aeppl/scan.py
@@ -0,0 +1,512 @@
+from copy import copy
+from typing import Callable, Dict, Iterable, List, Tuple, cast
+
+import aesara
+import aesara.tensor as at
+import numpy as np
+from aesara.graph.basic import Variable
+from aesara.graph.fg import FunctionGraph
+from aesara.graph.op import compute_test_value
+from aesara.graph.opt import local_optimizer, out2in
+from aesara.graph.optdb import OptimizationQuery
+from aesara.scan.op import Scan
+from aesara.scan.opt import scan_eqopt1, scan_eqopt2
+from aesara.scan.utils import ScanArgs
+from aesara.tensor.random.type import RandomType
+from aesara.tensor.subtensor import Subtensor, indices_from_subtensor
+from aesara.tensor.var import TensorVariable
+from aesara.updates import OrderedUpdates
+
+from aeppl.abstract import MeasurableVariable, _get_measurable_outputs
+from aeppl.joint_logprob import factorized_joint_logprob
+from aeppl.logprob import _logprob
+from aeppl.opt import inc_subtensor_ops, logprob_rewrites_db
+
+
+class MeasurableScan(Scan):
+    """A placeholder used to specify a log-likelihood for a scan sub-graph."""
+
+
+MeasurableVariable.register(MeasurableScan)
+
+
+def convert_outer_out_to_in(
+    input_scan_args: ScanArgs,
+    outer_out_vars: Iterable[TensorVariable],
+    new_outer_input_vars: Dict[TensorVariable, TensorVariable],
+    inner_out_fn: Callable[
+        [Dict[TensorVariable, TensorVariable]], Iterable[TensorVariable]
+    ],
+) -> ScanArgs:
+    r"""Convert outer-graph outputs into outer-graph inputs.
+
+    Parameters
+    ==========
+    input_scan_args:
+        The source `Scan` arguments.
+    outer_out_vars:
+        The outer-graph output variables that are to be converted into an
+        outer-graph input.
+    new_outer_input_vars:
+        The variables used for the new outer-graph input computed for
+        `outer_out_vars`.
+    inner_out_fn:
+        A function that takes the remapped outer-out variables and produces new
+        inner-graph outputs.  This can be used to transform the
+        `outer_out_vars`\s' corresponding inner-graph outputs into something
+        else entirely, like log-probabilities.
+
+    Outputs
+    =======
+    A `ScanArgs` object for a `Scan` in which `outer_out_vars` has been converted to an
+    outer-graph input.
+    """
+
+    output_scan_args = copy(input_scan_args)
+    inner_outs_to_new_inner_ins = {}
+
+    # Map inner-outputs to outer-outputs
+    old_inner_outs_to_outer_outs = {}
+
+    for oo_var in outer_out_vars:
+
+        var_info = output_scan_args.find_among_fields(
+            oo_var, field_filter=lambda x: x.startswith("outer_out")
+        )
+
+        assert var_info is not None
+        assert oo_var in new_outer_input_vars
+
+        io_var = output_scan_args.get_alt_field(var_info, "inner_out")
+        old_inner_outs_to_outer_outs[io_var] = oo_var
+
+    # In this loop, we gather information about the new inner-inputs that have
+    # been created and what their corresponding inner-outputs were, and we
+    # update the outer and inner-inputs to reflect the addition of new
+    # inner-inputs.
+    for old_inner_out_var, oo_var in old_inner_outs_to_outer_outs.items():
+
+        # Couldn't one do the same with `var_info`?
+        inner_out_info = output_scan_args.find_among_fields(
+            old_inner_out_var, field_filter=lambda x: x.startswith("inner_out")
+        )
+
+        output_scan_args.remove_from_fields(old_inner_out_var, rm_dependents=False)
+
+        # Remove the old outer-output variable.
+        # Not sure if this really matters, since we don't use the outer-outputs
+        # when building a new `Scan`, but doing it keeps the `ScanArgs` object
+        # consistent.
+        output_scan_args.remove_from_fields(oo_var, rm_dependents=False)
+
+        # Use the index for the specific inner-graph sub-collection to which this
+        # variable belongs (e.g. index `1` among the inner-graph sit-sot terms)
+        var_idx = inner_out_info.index
+
+        # The old inner-output variable becomes the a new inner-input
+        new_inner_in_var = old_inner_out_var.clone()
+        if new_inner_in_var.name:
+            new_inner_in_var.name = f"{new_inner_in_var.name}_vv"
+
+        inner_outs_to_new_inner_ins[old_inner_out_var] = new_inner_in_var
+
+        # We want to remove elements from both lists and tuples, because the
+        # members of `ScanArgs` could switch from being `list`s to `tuple`s
+        # soon
+        def remove(x, i):
+            return x[:i] + x[i + 1 :]
+
+        # If we're replacing a [m|s]it-sot, then we need to add a new nit-sot
+        add_nit_sot = False
+        if inner_out_info.name.endswith("mit_sot"):
+            inner_in_mit_sot_var = cast(
+                Tuple[int, ...], tuple(output_scan_args.inner_in_mit_sot[var_idx])
+            )
+            new_inner_in_seqs = inner_in_mit_sot_var + (new_inner_in_var,)
+            new_inner_in_mit_sot = remove(output_scan_args.inner_in_mit_sot, var_idx)
+            new_outer_in_mit_sot = remove(output_scan_args.outer_in_mit_sot, var_idx)
+            new_inner_in_sit_sot = tuple(output_scan_args.inner_in_sit_sot)
+            new_outer_in_sit_sot = tuple(output_scan_args.outer_in_sit_sot)
+            add_nit_sot = True
+        elif inner_out_info.name.endswith("sit_sot"):
+            new_inner_in_seqs = (output_scan_args.inner_in_sit_sot[var_idx],) + (
+                new_inner_in_var,
+            )
+            new_inner_in_sit_sot = remove(output_scan_args.inner_in_sit_sot, var_idx)
+            new_outer_in_sit_sot = remove(output_scan_args.outer_in_sit_sot, var_idx)
+            new_inner_in_mit_sot = tuple(output_scan_args.inner_in_mit_sot)
+            new_outer_in_mit_sot = tuple(output_scan_args.outer_in_mit_sot)
+            add_nit_sot = True
+        else:
+            new_inner_in_seqs = (new_inner_in_var,)
+            new_inner_in_mit_sot = tuple(output_scan_args.inner_in_mit_sot)
+            new_outer_in_mit_sot = tuple(output_scan_args.outer_in_mit_sot)
+            new_inner_in_sit_sot = tuple(output_scan_args.inner_in_sit_sot)
+            new_outer_in_sit_sot = tuple(output_scan_args.outer_in_sit_sot)
+
+        output_scan_args.inner_in_mit_sot = list(new_inner_in_mit_sot)
+        output_scan_args.inner_in_sit_sot = list(new_inner_in_sit_sot)
+        output_scan_args.outer_in_mit_sot = list(new_outer_in_mit_sot)
+        output_scan_args.outer_in_sit_sot = list(new_outer_in_sit_sot)
+
+        if inner_out_info.name.endswith("mit_sot"):
+            mit_sot_var_taps = cast(
+                Tuple[int, ...], tuple(output_scan_args.mit_sot_in_slices[var_idx])
+            )
+            taps = mit_sot_var_taps + (0,)
+            new_mit_sot_in_slices = remove(output_scan_args.mit_sot_in_slices, var_idx)
+        elif inner_out_info.name.endswith("sit_sot"):
+            taps = (-1, 0)
+            new_mit_sot_in_slices = tuple(output_scan_args.mit_sot_in_slices)
+        else:
+            taps = (0,)
+            new_mit_sot_in_slices = tuple(output_scan_args.mit_sot_in_slices)
+
+        output_scan_args.mit_sot_in_slices = list(new_mit_sot_in_slices)
+
+        taps, new_inner_in_seqs = zip(
+            *sorted(zip(taps, new_inner_in_seqs), key=lambda x: x[0])
+        )
+
+        new_inner_in_seqs = tuple(output_scan_args.inner_in_seqs) + tuple(
+            reversed(new_inner_in_seqs)
+        )
+
+        output_scan_args.inner_in_seqs = list(new_inner_in_seqs)
+
+        slice_seqs = zip(
+            -np.asarray(taps), [n if n < 0 else None for n in reversed(taps)]
+        )
+
+        # XXX: If the caller passes the variables output by `aesara.scan`, it's
+        # likely that this will fail, because those variables can sometimes be
+        # slices of the actual outer-inputs (e.g. `out[1:]` instead of `out`
+        # when `taps=[-1]`).
+        var_slices = [new_outer_input_vars[oo_var][b:e] for b, e in slice_seqs]
+        n_steps = at.min([at.shape(n)[0] for n in var_slices])
+
+        output_scan_args.n_steps = n_steps
+
+        new_outer_in_seqs = tuple(output_scan_args.outer_in_seqs) + tuple(
+            v[:n_steps] for v in var_slices
+        )
+
+        output_scan_args.outer_in_seqs = list(new_outer_in_seqs)
+
+        if add_nit_sot:
+            new_outer_in_nit_sot = tuple(output_scan_args.outer_in_nit_sot) + (n_steps,)
+        else:
+            new_outer_in_nit_sot = tuple(output_scan_args.outer_in_nit_sot)
+
+        output_scan_args.outer_in_nit_sot = list(new_outer_in_nit_sot)
+
+    # Now, we can add new inner-outputs for the custom calculations.
+    # We don't need to create corresponding outer-outputs, because `Scan` will
+    # do that when we call `Scan.make_node`.  All we need is a consistent
+    # outer-inputs and inner-graph spec., which we should have in
+    # `output_scan_args`.
+    remapped_io_to_ii = inner_outs_to_new_inner_ins
+    new_inner_out_nit_sot = tuple(output_scan_args.inner_out_nit_sot) + tuple(
+        inner_out_fn(remapped_io_to_ii)
+    )
+
+    output_scan_args.inner_out_nit_sot = list(new_inner_out_nit_sot)
+
+    return output_scan_args
+
+
+def get_random_outer_outputs(
+    scan_args: ScanArgs,
+) -> List[Tuple[int, TensorVariable, TensorVariable]]:
+    """Get the `MeasurableVariable` outputs of a `Scan` (well, its `ScanArgs`).
+
+    Returns
+    =======
+    A tuple of tuples containing the index of each outer-output variable, the
+    outer-output variable itself, and the inner-output variable that
+    is an instance of `MeasurableVariable`.
+    """
+    rv_vars = []
+    for n, oo_var in enumerate(
+        [o for o in scan_args.outer_outputs if not isinstance(o.type, RandomType)]
+    ):
+        oo_info = scan_args.find_among_fields(oo_var)
+        io_type = oo_info.name[(oo_info.name.index("_", 6) + 1) :]
+        inner_out_type = "inner_out_{}".format(io_type)
+        io_var = getattr(scan_args, inner_out_type)[oo_info.index]
+        if io_var.owner and isinstance(io_var.owner.op, MeasurableVariable):
+            rv_vars.append((n, oo_var, io_var))
+    return rv_vars
+
+
+def construct_scan(scan_args: ScanArgs) -> Tuple[List[TensorVariable], OrderedUpdates]:
+    scan_op = Scan(scan_args.inner_inputs, scan_args.inner_outputs, scan_args.info)
+    node = scan_op.make_node(*scan_args.outer_inputs)
+    updates = OrderedUpdates(zip(scan_args.outer_in_shared, scan_args.outer_out_shared))
+    return node.outputs, updates
+
+
+@_logprob.register(MeasurableScan)
+def logprob_ScanRV(op, values, *inputs, name=None, **kwargs):
+
+    new_node = op.make_node(*inputs)
+    scan_args = ScanArgs.from_node(new_node)
+    rv_outer_outs = get_random_outer_outputs(scan_args)
+
+    var_indices, rv_vars, io_vars = zip(*rv_outer_outs)
+    value_map = {_rv: _val for _rv, _val in zip(rv_vars, values)}
+
+    def create_inner_out_logp(
+        value_map: Dict[TensorVariable, TensorVariable]
+    ) -> TensorVariable:
+        """Create a log-likelihood inner-output for a `Scan`."""
+        logp_parts = factorized_joint_logprob(value_map, warn_missing_rvs=False)
+        return logp_parts.values()
+
+    logp_scan_args = convert_outer_out_to_in(
+        scan_args,
+        rv_vars,
+        value_map,
+        inner_out_fn=create_inner_out_logp,
+    )
+
+    # Remove the shared variables corresponding to replaced terms.
+
+    # TODO FIXME: This is a really dirty approach, because it effectively
+    # assumes that all sampling is being removed, and, thus, all shared updates
+    # relating to `RandomType`s.  Instead, we should be more precise and only
+    # remove the `RandomType`s associated with `values`.
+    logp_scan_args.outer_in_shared = [
+        i for i in logp_scan_args.outer_in_shared if not isinstance(i.type, RandomType)
+    ]
+    logp_scan_args.inner_in_shared = [
+        i for i in logp_scan_args.inner_in_shared if not isinstance(i.type, RandomType)
+    ]
+    logp_scan_args.inner_out_shared = [
+        i for i in logp_scan_args.inner_out_shared if not isinstance(i.type, RandomType)
+    ]
+    # XXX TODO: Remove this properly
+    # logp_scan_args.outer_out_shared = []
+
+    logp_scan_out, updates = construct_scan(logp_scan_args)
+
+    # Automatically pick up updates so that we don't have to pass them around
+    for key, value in updates.items():
+        key.default_update = value
+
+    return logp_scan_out
+
+
+@local_optimizer([Scan])
+def find_measurable_scans(fgraph, node):
+    r"""Finds `Scan`\s for which a `logprob` can be computed.
+
+    This will convert said `Scan`\s into `MeasurableScan`\s.  It also updates
+    random variable and value variable mappings that have been specified for
+    parts of a `Scan`\s outputs (e.g. everything except the initial values).
+    """
+
+    if not isinstance(node.op, Scan):
+        return None
+
+    if isinstance(node.op, MeasurableScan):
+        return None
+
+    if not hasattr(fgraph, "shape_feature"):
+        return None  # pragma: no cover
+
+    rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
+
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    curr_scanargs = ScanArgs.from_node(node)
+
+    # Find the un-output `MeasurableVariable`s created in the inner-graph
+    clients: Dict[Variable, List[Variable]] = {}
+
+    local_fgraph_topo = aesara.graph.basic.io_toposort(
+        curr_scanargs.inner_inputs,
+        [o for o in curr_scanargs.inner_outputs if not isinstance(o.type, RandomType)],
+        clients=clients,
+    )
+    for n in local_fgraph_topo:
+        if isinstance(n.op, MeasurableVariable):
+            non_output_node_clients = [
+                c for c in clients[n] if c not in curr_scanargs.inner_outputs
+            ]
+
+            if non_output_node_clients:
+                # This node is a `MeasurableVariable`, but it depends on
+                # variable that's not being output?
+                # TODO: Why can't we make this a `MeasurableScan`?
+                return None
+
+    if not any(out in rv_map_feature.rv_values for out in node.outputs):
+        # We need to remap user inputs that have been specified in terms of
+        # `Subtensor`s of this `Scan`'s node's outputs.
+        #
+        # For example, the output that the user got was something like
+        # `out[1:]` for `outputs_info = [{"initial": x0, "taps": [-1]}]`, so
+        # they likely passed `{out[1:]: x_1T_vv}` to `joint_logprob`.
+        # Since `out[1:]` isn't really the output of a `Scan`, but a
+        # `Subtensor` of the output `out` of a `Scan`, we need to account for
+        # that.
+
+        # Get any `Subtensor` outputs that have been applied to outputs of this
+        # `Scan` (and get the corresponding indices of the outputs from this
+        # `Scan`)
+        output_clients: List[Tuple[Variable, int]] = sum(
+            [
+                [
+                    # This is expected to work for `Subtensor` `Op`s,
+                    # because they only ever have one output
+                    (cl.default_output(), i)
+                    for cl, _ in fgraph.get_clients(out)
+                    if isinstance(cl.op, Subtensor)
+                ]
+                for i, out in enumerate(node.outputs)
+            ],
+            [],
+        )
+
+        # The second items in these tuples are the value variables mapped to
+        # the *user-specified* measurable variables (i.e. the first items) that
+        # are `Subtensor`s of the outputs of this `Scan`.  The second items are
+        # the index of the corresponding output of this `Scan` node.
+        indirect_rv_vars = [
+            (out, rv_map_feature.rv_values[out], out_idx)
+            for out, out_idx in output_clients
+            if out in rv_map_feature.rv_values
+        ]
+
+        if not indirect_rv_vars:
+            return None
+
+        # We need this for the `clone` in the loop that follows
+        if aesara.config.compute_test_value != "off":
+            compute_test_value(node)
+
+        # We're going to replace the user's random variable/value variable mappings
+        # with ones that map directly to outputs of this `Scan`.
+        for rv_var, val_var, out_idx in indirect_rv_vars:
+
+            # The full/un-`Subtensor`ed `Scan` output that we need to use
+            full_out = node.outputs[out_idx]
+
+            assert rv_var.owner.inputs[0] == full_out
+
+            # A new value variable that spans the full output.
+            # We don't want the old graph to appear in the new log-probability
+            # graph, so we use the shape feature to (hopefully) get the shape
+            # without the entire `Scan` itself.
+            full_out_shape = tuple(
+                fgraph.shape_feature.get_shape(full_out, i)
+                for i in range(full_out.ndim)
+            )
+            new_val_var = at.empty(full_out_shape, dtype=full_out.dtype)
+
+            # Set the parts of this new value variable that applied to the
+            # user-specified value variable to the user's value variable
+            subtensor_indices = indices_from_subtensor(
+                rv_var.owner.inputs[1:], rv_var.owner.op.idx_list
+            )
+            # E.g. for a single `-1` TAPS, `s_0T[1:] = s_1T` where `s_0T` is
+            # `new_val_var` and `s_1T` is the user-specified value variable
+            # that only spans times `t=1` to `t=T`.
+            new_val_var = at.set_subtensor(new_val_var[subtensor_indices], val_var)
+
+            # This is the outer-input that sets `s_0T[i] = taps[i]` where `i`
+            # is a TAP index (e.g. a TAP of `-1` maps to index `0` in a vector
+            # of the entire series).
+            var_info = curr_scanargs.find_among_fields(full_out)
+            alt_type = var_info.name[(var_info.name.index("_", 6) + 1) :]
+            outer_input_var = getattr(curr_scanargs, f"outer_in_{alt_type}")[
+                var_info.index
+            ]
+
+            # These outer-inputs are using by `aesara.scan.utils.expand_empty`, and
+            # are expected to consist of only a single `set_subtensor` call.
+            # That's why we can simply replace the first argument of the node.
+            assert isinstance(outer_input_var.owner.op, inc_subtensor_ops)
+
+            # We're going to set those values on our `new_val_var` so that it can
+            # serve as a complete replacement for the old input `outer_input_var`.
+            # from aesara.graph import clone_replace
+            #
+            new_val_var = outer_input_var.owner.clone_with_new_inputs(
+                [new_val_var] + outer_input_var.owner.inputs[1:]
+            ).default_output()
+
+            # Replace the mapping
+            del rv_map_feature.rv_values[rv_var]
+            rv_map_feature.rv_values[full_out] = new_val_var
+
+    op = MeasurableScan(
+        curr_scanargs.inner_inputs, curr_scanargs.inner_outputs, curr_scanargs.info
+    )
+    new_node = op.make_node(*curr_scanargs.outer_inputs)
+
+    return dict(zip(node.outputs, new_node.outputs))
+
+
+@local_optimizer([Scan])
+def add_opts_to_inner_graphs(fgraph, node):
+    """Update the `Mode`(s) used to compile the inner-graph of a `Scan` `Op`.
+
+    This is how we add `rv_sinking_db` and any other log-probability rewrites
+    to the "body" (i.e. inner-graph) of a `Scan` loop.
+    """
+
+    if not isinstance(node.op, Scan):
+        return None
+
+    inner_fgraph = FunctionGraph(
+        node.op.inputs,
+        node.op.outputs,
+        clone=True,
+        copy_inputs=False,
+        copy_orphans=False,
+    )
+
+    logprob_rewrites_db.query(OptimizationQuery(include=["basic"])).optimize(
+        inner_fgraph
+    )
+
+    new_outputs = list(inner_fgraph.outputs)
+
+    op = Scan(node.op.inputs, new_outputs, node.op.info)
+    new_node = op.make_node(*node.inputs)
+
+    return dict(zip(node.outputs, new_node.outputs))
+
+
+@_get_measurable_outputs.register(MeasurableScan)
+def _get_measurable_outputs_MeasurableScan(op, node):
+    # TODO: This should probably use `get_random_outer_outputs`
+    # scan_args = ScanArgs.from_node(node)
+    # rv_outer_outs = get_random_outer_outputs(scan_args)
+    return [o for o in node.outputs if not isinstance(o.type, RandomType)]
+
+
+logprob_rewrites_db.register(
+    "add_opts_to_inner_graphs",
+    out2in(
+        add_opts_to_inner_graphs, name="add_opts_to_inner_graphs", ignore_newtrees=True
+    ),
+    -100,
+    "basic",
+    "scan",
+)
+
+logprob_rewrites_db.register(
+    "find_measurable_scans",
+    out2in(find_measurable_scans, name="find_measurable_scans", ignore_newtrees=True),
+    0,
+    "basic",
+    "scan",
+)
+
+# Add scan canonicalizations that aren't in the canonicalize DB
+logprob_rewrites_db.register("scan_eqopt1", scan_eqopt1, -9, "basic", "scan")
+logprob_rewrites_db.register("scan_eqopt2", scan_eqopt2, -9, "basic", "scan")
diff --git a/setup.py b/setup.py
index a85b2fcc..2e51755f 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,7 @@
     install_requires=[
         "numpy>=1.18.1",
         "scipy>=1.4.0",
-        "aesara >= 2.1.1",
+        "aesara >= 2.2.1",
     ],
     tests_require=["pytest"],
     long_description=open("README.rst").read() if exists("README.rst") else "",


Test patch:
diff --git a/tests/test_scan.py b/tests/test_scan.py
new file mode 100644
index 00000000..9c6b0367
--- /dev/null
+++ b/tests/test_scan.py
@@ -0,0 +1,393 @@
+import aesara
+import aesara.tensor as at
+import numpy as np
+import pytest
+from aesara.scan.utils import ScanArgs
+
+from aeppl.joint_logprob import factorized_joint_logprob, joint_logprob
+from aeppl.logprob import logprob
+from aeppl.scan import construct_scan, convert_outer_out_to_in, get_random_outer_outputs
+from tests.utils import assert_no_rvs
+
+
+def create_inner_out_logp(value_map):
+    """Create a log-likelihood inner-output.
+
+    This is intended to be use with `get_random_outer_outputs`.
+
+    """
+    res = []
+    for old_inner_out_var, new_inner_in_var in value_map.items():
+        logp = logprob(old_inner_out_var, new_inner_in_var)
+        if new_inner_in_var.name:
+            logp.name = "logp({})".format(new_inner_in_var.name)
+        res.append(logp)
+
+    return res
+
+
+def test_convert_outer_out_to_in_sit_sot():
+    """Test a single replacement with `convert_outer_out_to_in`.
+
+    This should be a single SIT-SOT replacement.
+    """
+
+    rng_state = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(1234)))
+    rng_tt = aesara.shared(rng_state, name="rng", borrow=True)
+    rng_tt.tag.is_rng = True
+    rng_tt.default_update = rng_tt
+
+    #
+    # We create a `Scan` representing a time-series model with normally
+    # distributed responses that are dependent on lagged values of both the
+    # response `RandomVariable` and a lagged "deterministic" that also depends
+    # on the lagged response values.
+    #
+    def input_step_fn(mu_tm1, y_tm1, rng):
+        mu_tm1.name = "mu_tm1"
+        y_tm1.name = "y_tm1"
+        mu = mu_tm1 + y_tm1 + 1
+        mu.name = "mu_t"
+        return mu, at.random.normal(mu, 1.0, rng=rng, name="Y_t")
+
+    (mu_tt, Y_rv), _ = aesara.scan(
+        fn=input_step_fn,
+        outputs_info=[
+            {
+                "initial": at.as_tensor_variable(0.0, dtype=aesara.config.floatX),
+                "taps": [-1],
+            },
+            {
+                "initial": at.as_tensor_variable(0.0, dtype=aesara.config.floatX),
+                "taps": [-1],
+            },
+        ],
+        non_sequences=[rng_tt],
+        n_steps=10,
+    )
+
+    mu_tt.name = "mu_tt"
+    mu_tt.owner.inputs[0].name = "mu_all"
+    Y_rv.name = "Y_rv"
+    Y_all = Y_rv.owner.inputs[0]
+    Y_all.name = "Y_all"
+
+    input_scan_args = ScanArgs.from_node(Y_rv.owner.inputs[0].owner)
+
+    # TODO FIXME: Everything below needs to be replaced with explicit asserts
+    # on the values in `input_scan_args`
+
+    #
+    # Sample from the model and create another `Scan` that computes the
+    # log-likelihood of the model at the sampled point.
+    #
+    Y_obs = at.as_tensor_variable(Y_rv.eval())
+    Y_obs.name = "Y_obs"
+
+    def output_step_fn(y_t, y_tm1, mu_tm1):
+        mu_tm1.name = "mu_tm1"
+        y_tm1.name = "y_tm1"
+        mu = mu_tm1 + y_tm1 + 1
+        mu.name = "mu_t"
+        logp = logprob(at.random.normal(mu, 1.0), y_t)
+        logp.name = "logp"
+        return mu, logp
+
+    (mu_tt, Y_logp), _ = aesara.scan(
+        fn=output_step_fn,
+        sequences=[{"input": Y_obs, "taps": [0, -1]}],
+        outputs_info=[
+            {
+                "initial": at.as_tensor_variable(0.0, dtype=aesara.config.floatX),
+                "taps": [-1],
+            },
+            {},
+        ],
+    )
+
+    Y_logp.name = "Y_logp"
+    mu_tt.name = "mu_tt"
+
+    #
+    # Get the model output variable that corresponds to the response
+    # `RandomVariable`
+    #
+    oo_idx, oo_var, io_var = get_random_outer_outputs(input_scan_args)[0]
+
+    #
+    # Convert the original model `Scan` into another `Scan` that's equivalent
+    # to the log-likelihood `Scan` given above.
+    # In other words, automatically construct the log-likelihood `Scan` based
+    # on the model `Scan`.
+    #
+    value_map = {Y_all: Y_obs}
+    test_scan_args = convert_outer_out_to_in(
+        input_scan_args,
+        [oo_var],
+        value_map,
+        inner_out_fn=create_inner_out_logp,
+    )
+
+    scan_out, updates = construct_scan(test_scan_args)
+
+    #
+    # Evaluate the manually and automatically constructed log-likelihoods and
+    # compare.
+    #
+    res = scan_out[oo_idx].eval()
+    exp_res = Y_logp.eval()
+
+    assert np.array_equal(res, exp_res)
+
+
+def test_convert_outer_out_to_in_mit_sot():
+    """Test a single replacement with `convert_outer_out_to_in`.
+
+    This should be a single MIT-SOT replacement.
+    """
+
+    rng_state = np.random.default_rng(1234)
+    rng_tt = aesara.shared(rng_state, name="rng", borrow=True)
+    rng_tt.tag.is_rng = True
+    rng_tt.default_update = rng_tt
+
+    #
+    # This is a very simple model with only one output, but multiple
+    # taps/lags.
+    #
+    def input_step_fn(y_tm1, y_tm2, rng):
+        y_tm1.name = "y_tm1"
+        y_tm2.name = "y_tm2"
+        return at.random.normal(y_tm1 + y_tm2, 1.0, rng=rng, name="Y_t")
+
+    Y_rv, _ = aesara.scan(
+        fn=input_step_fn,
+        outputs_info=[
+            {"initial": at.as_tensor_variable(np.r_[-1.0, 0.0]), "taps": [-1, -2]},
+        ],
+        non_sequences=[rng_tt],
+        n_steps=10,
+    )
+
+    Y_rv.name = "Y_rv"
+    Y_all = Y_rv.owner.inputs[0]
+    Y_all.name = "Y_all"
+
+    Y_obs = at.as_tensor_variable(Y_rv.eval())
+    Y_obs.name = "Y_obs"
+
+    input_scan_args = ScanArgs.from_node(Y_rv.owner.inputs[0].owner)
+
+    # TODO FIXME: Everything below needs to be replaced with explicit asserts
+    # on the values in `input_scan_args`
+
+    #
+    # The corresponding log-likelihood
+    #
+    def output_step_fn(y_t, y_tm1, y_tm2):
+        y_t.name = "y_t"
+        y_tm1.name = "y_tm1"
+        y_tm2.name = "y_tm2"
+        logp = logprob(at.random.normal(y_tm1 + y_tm2, 1.0), y_t)
+        logp.name = "logp(y_t)"
+        return logp
+
+    Y_logp, _ = aesara.scan(
+        fn=output_step_fn,
+        sequences=[{"input": Y_obs, "taps": [0, -1, -2]}],
+        outputs_info=[{}],
+    )
+
+    #
+    # Get the model output variable that corresponds to the response
+    # `RandomVariable`
+    #
+    oo_idx, oo_var, io_var = get_random_outer_outputs(input_scan_args)[0]
+
+    #
+    # Convert the original model `Scan` into another `Scan` that's equivalent
+    # to the log-likelihood `Scan` given above.
+    # In other words, automatically construct the log-likelihood `Scan` based
+    # on the model `Scan`.
+
+    value_map = {Y_all: Y_obs}
+    test_scan_args = convert_outer_out_to_in(
+        input_scan_args,
+        [oo_var],
+        value_map,
+        inner_out_fn=create_inner_out_logp,
+    )
+
+    scan_out, updates = construct_scan(test_scan_args)
+
+    #
+    # Evaluate the manually and automatically constructed log-likelihoods and
+    # compare.
+    #
+    res = scan_out[oo_idx].eval()
+    exp_res = Y_logp.eval()
+
+    assert np.array_equal(res, exp_res)
+
+
+@pytest.mark.parametrize(
+    "require_inner_rewrites",
+    [
+        False,
+        True,
+    ],
+)
+def test_scan_joint_logprob(require_inner_rewrites):
+    srng = at.random.RandomStream()
+
+    N_tt = at.iscalar("N")
+    N_val = 10
+    N_tt.tag.test_value = N_val
+
+    M_tt = at.iscalar("M")
+    M_val = 2
+    M_tt.tag.test_value = M_val
+
+    mus_tt = at.matrix("mus_t")
+
+    mus_val = np.stack([np.arange(0.0, 10), np.arange(0.0, -10, -1)], axis=-1).astype(
+        aesara.config.floatX
+    )
+    mus_tt.tag.test_value = mus_val
+
+    sigmas_tt = at.ones((N_tt,))
+    Gamma_rv = srng.dirichlet(at.ones((M_tt, M_tt)), name="Gamma")
+
+    Gamma_vv = Gamma_rv.clone()
+    Gamma_vv.name = "Gamma_vv"
+
+    Gamma_val = np.array([[0.5, 0.5], [0.5, 0.5]])
+    Gamma_rv.tag.test_value = Gamma_val
+
+    def scan_fn(mus_t, sigma_t, Gamma_t):
+        S_t = srng.categorical(Gamma_t[0], name="S_t")
+
+        if require_inner_rewrites:
+            Y_t = srng.normal(mus_t, sigma_t, name="Y_t")[S_t]
+        else:
+            Y_t = srng.normal(mus_t[S_t], sigma_t, name="Y_t")
+
+        return Y_t, S_t
+
+    (Y_rv, S_rv), _ = aesara.scan(
+        fn=scan_fn,
+        sequences=[mus_tt, sigmas_tt],
+        non_sequences=[Gamma_rv],
+        outputs_info=[{}, {}],
+        strict=True,
+        name="scan_rv",
+    )
+    Y_rv.name = "Y"
+    S_rv.name = "S"
+
+    y_vv = Y_rv.clone()
+    y_vv.name = "y"
+
+    s_vv = S_rv.clone()
+    s_vv.name = "s"
+
+    y_logp = joint_logprob({Y_rv: y_vv, S_rv: s_vv, Gamma_rv: Gamma_vv})
+
+    y_val = np.arange(10)
+    s_val = np.array([0, 1, 0, 1, 1, 0, 0, 0, 1, 1])
+
+    test_point = {
+        y_vv: y_val,
+        s_vv: s_val,
+        M_tt: M_val,
+        N_tt: N_val,
+        mus_tt: mus_val,
+        Gamma_vv: Gamma_val,
+    }
+
+    y_logp_fn = aesara.function(list(test_point.keys()), y_logp)
+
+    assert_no_rvs(y_logp_fn.maker.fgraph.outputs[0])
+
+    # Construct the joint log-probability by hand so we can compare it with
+    # `y_logp`
+    def scan_fn(mus_t, sigma_t, Y_t_val, S_t_val, Gamma_t):
+        S_t = at.random.categorical(Gamma_t[0], name="S_t")
+        Y_t = at.random.normal(mus_t[S_t_val], sigma_t, name="Y_t")
+        Y_t_logp, S_t_logp = logprob(Y_t, Y_t_val), logprob(S_t, S_t_val)
+        Y_t_logp.name = "log(Y_t=y_t)"
+        S_t_logp.name = "log(S_t=s_t)"
+        return Y_t_logp, S_t_logp
+
+    (Y_rv_logp, S_rv_logp), _ = aesara.scan(
+        fn=scan_fn,
+        sequences=[mus_tt, sigmas_tt, y_vv, s_vv],
+        non_sequences=[Gamma_vv],
+        outputs_info=[{}, {}],
+        strict=True,
+        name="scan_rv",
+    )
+    Y_rv_logp.name = "logp(Y=y)"
+    S_rv_logp.name = "logp(S=s)"
+
+    Gamma_logp = logprob(Gamma_rv, Gamma_vv)
+
+    y_logp_ref = Y_rv_logp.sum() + S_rv_logp.sum() + Gamma_logp.sum()
+
+    assert_no_rvs(y_logp_ref)
+
+    y_logp_val = y_logp.eval(test_point)
+
+    y_logp_ref_val = y_logp_ref.eval(test_point)
+
+    assert np.allclose(y_logp_val, y_logp_ref_val)
+
+
+@aesara.config.change_flags(compute_test_value="raise")
+def test_initial_values():
+    srng = at.random.RandomStream(seed=2320)
+
+    p_S_0 = np.array([0.9, 0.1])
+    S_0_rv = srng.categorical(p_S_0, name="S_0")
+    S_0_rv.tag.test_value = 0
+
+    Gamma_at = at.matrix("Gamma")
+    Gamma_at.tag.test_value = np.array([[0, 1], [1, 0]])
+
+    s_0_vv = S_0_rv.clone()
+    s_0_vv.name = "s_0"
+
+    def step_fn(S_tm1, Gamma):
+        S_t = srng.categorical(Gamma[S_tm1], name="S_t")
+        return S_t
+
+    S_1T_rv, _ = aesara.scan(
+        fn=step_fn,
+        outputs_info=[{"initial": S_0_rv, "taps": [-1]}],
+        non_sequences=[Gamma_at],
+        strict=True,
+        n_steps=10,
+        name="S_0T",
+    )
+
+    S_1T_rv.name = "S_1T"
+    s_1T_vv = S_1T_rv.clone()
+    s_1T_vv.name = "s_1T"
+
+    logp_parts = factorized_joint_logprob({S_1T_rv: s_1T_vv, S_0_rv: s_0_vv})
+
+    s_0_val = 0
+    s_1T_val = np.array([1, 0, 1, 0, 1, 1, 0, 1, 0, 1])
+    Gamma_val = np.array([[0.1, 0.9], [0.9, 0.1]])
+
+    exp_res = np.log(p_S_0[s_0_val])
+    s_prev = s_0_val
+    for s in s_1T_val:
+        exp_res += np.log(Gamma_val[s_prev, s])
+        s_prev = s
+
+    S_0T_logp = sum(v.sum() for v in logp_parts.values())
+    S_0T_logp_fn = aesara.function([s_0_vv, s_1T_vv, Gamma_at], S_0T_logp)
+    res = S_0T_logp_fn(s_0_val, s_1T_val, Gamma_val)
+
+    assert res == pytest.approx(exp_res)
diff --git a/tests/utils.py b/tests/utils.py
index 07ba6cda..26a78034 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,9 +1,24 @@
-from aesara.graph.basic import ancestors
-from aesara.tensor.random.op import RandomVariable
+from aesara.graph.basic import walk
+from aesara.scan.op import Scan
+
+from aeppl.abstract import MeasurableVariable
 
 
 def assert_no_rvs(var):
-    assert not any(
-        isinstance(v.owner.op, RandomVariable) for v in ancestors([var]) if v.owner
-    )
-    return var
+    """Assert that there are no `MeasurableVariable` nodes in a graph."""
+
+    def expand(r):
+        owner = r.owner
+        if owner:
+            inputs = list(reversed(owner.inputs))
+
+            # TODO: We need a better--potentially type-based--means of
+            # determining whether or not an inner-graph is present
+            if isinstance(owner.op, Scan):
+                inputs += owner.op.outputs
+
+            return inputs
+
+    for v in walk([var], expand, False):
+        if v.owner and isinstance(v.owner.op, MeasurableVariable):
+            raise AssertionError(f"Variable {v} is a MeasurableVariable")