Fix logprob inference for scans with carried deterministic states

Author: ricardoV94

pymc/logprob/scan.py

@@ -42,15 +42,13 @@
 import pytensor.tensor as pt
 
 from pytensor.graph.basic import Variable
-from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import compute_test_value
 from pytensor.graph.rewriting.basic import node_rewriter
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
 from pytensor.scan.op import Scan
 from pytensor.scan.rewriting import scan_eqopt1, scan_eqopt2
 from pytensor.scan.utils import ScanArgs
 from pytensor.tensor.random.type import RandomType
-from pytensor.tensor.rewriting.shape import ShapeFeature
 from pytensor.tensor.subtensor import Subtensor, indices_from_subtensor
 from pytensor.tensor.var import TensorVariable
 from pytensor.updates import OrderedUpdates
@@ -63,11 +61,12 @@
 )
 from pymc.logprob.joint_logprob import factorized_joint_logprob
 from pymc.logprob.rewriting import (
-    PreserveRVMappings,
+    construct_ir_fgraph,
     inc_subtensor_ops,
     logprob_rewrites_db,
     measurable_ir_rewrites_db,
 )
+from pymc.pytensorf import replace_rvs_by_values
 
 
 class MeasurableScan(Scan):
@@ -249,9 +248,27 @@ def remove(x, i):
     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)
 
+    # Finally, we need to replace any lingering references to the new
+    # internal variables that could be in the recurrent states needed
+    # to compute the new nit_sots
+    traced_outs = (
+        output_scan_args.inner_out_mit_sot
+        + output_scan_args.inner_out_sit_sot
+        + output_scan_args.inner_out_nit_sot
+    )
+    traced_outs = replace_rvs_by_values(traced_outs, rvs_to_values=remapped_io_to_ii)
+    # Update output mappings
+    n_mit_sot = len(output_scan_args.inner_out_mit_sot)
+    output_scan_args.inner_out_mit_sot = traced_outs[:n_mit_sot]
+    offset = n_mit_sot
+    n_sit_sot = len(output_scan_args.inner_out_sit_sot)
+    output_scan_args.inner_out_sit_sot = traced_outs[offset : offset + n_sit_sot]
+    offset += n_sit_sot
+    n_nit_sot = len(output_scan_args.inner_out_nit_sot)
+    output_scan_args.inner_out_nit_sot = traced_outs[offset : offset + n_nit_sot]
+
     return output_scan_args
 
 
@@ -331,7 +348,10 @@ def create_inner_out_logp(value_map: Dict[TensorVariable, TensorVariable]) -> Te
     for key, value in updates.items():
         key.default_update = value
 
-    return logp_scan_out
+    # Return only the logp outputs, not any potentially carried states
+    logp_outputs = logp_scan_out[-len(values) :]
+
+    return logp_outputs
 
 
 @node_rewriter([Scan])
@@ -504,19 +524,9 @@ def add_opts_to_inner_graphs(fgraph, node):
     if getattr(node.op.mode, "had_logprob_rewrites", False):
         return None
 
-    inner_fgraph = FunctionGraph(
-        node.op.inner_inputs,
-        node.op.inner_outputs,
-        clone=True,
-        copy_inputs=False,
-        copy_orphans=False,
-        features=[
-            ShapeFeature(),
-            PreserveRVMappings({}),
-        ],
-    )
-
-    logprob_rewrites_db.query(RewriteDatabaseQuery(include=["basic"])).rewrite(inner_fgraph)
+    inner_rv_values = {out: out.type() for out in node.op.inner_outputs}
+    ir_rewriter = logprob_rewrites_db.query(RewriteDatabaseQuery(include=["basic"]))
+    inner_fgraph, rv_values, _ = construct_ir_fgraph(inner_rv_values, ir_rewriter=ir_rewriter)
 
     new_outputs = list(inner_fgraph.outputs)
 
@@ -531,11 +541,23 @@ def add_opts_to_inner_graphs(fgraph, node):
 
 
 @_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)]
+def _get_measurable_outputs_MeasurableScan(op: Scan, node):
+    """Collect measurable outputs for Measurable Scans"""
+    inner_out_from_outer_out_map = op.get_oinp_iinp_iout_oout_mappings()["inner_out_from_outer_out"]
+    inner_outs = op.inner_outputs
+
+    # Measurable scan outputs are those whose inner scan output counterparts are also measurable
+    measurable_outputs = []
+    for out_idx, out in enumerate(node.outputs):
+        [inner_out_idx] = inner_out_from_outer_out_map[out_idx]
+        inner_out = inner_outs[inner_out_idx]
+        inner_out_node = inner_out.owner
+        if isinstance(
+            inner_out_node.op, MeasurableVariable
+        ) and inner_out in get_measurable_outputs(inner_out_node.op, inner_out_node):
+            measurable_outputs.append(out)
+
+    return measurable_outputs
 
 
 measurable_ir_rewrites_db.register(

tests/logprob/test_scan.py

@@ -502,3 +502,55 @@ def test_scan_over_seqs():
         ys_logp.eval({xs_vv: xs_test, ys_vv: ys_test}),
         stats.norm.logpdf(ys_test, xs_test),
     )
+
+
+def test_scan_carried_deterministic_state():
+    """Test logp of scans with carried states downstream of measured variables.
+
+    A moving average model with 2 lags is used for testing.
+    """
+    rng = np.random.default_rng(490)
+    steps = 99
+
+    rho = pt.vector("rho", shape=(2,))
+    sigma = pt.scalar("sigma")
+
+    def ma2_step(eps_tm2, eps_tm1, rho, sigma):
+        mu = eps_tm1 * rho[0] + eps_tm2 * rho[1]
+        y = pt.random.normal(mu, sigma)
+        eps = y - mu
+        update = {y.owner.inputs[0]: y.owner.outputs[0]}
+        return (eps, y), update
+
+    [_, ma2], ma2_updates = pytensor.scan(
+        fn=ma2_step,
+        outputs_info=[{"initial": pt.arange(2, dtype="float64"), "taps": range(-2, 0)}, None],
+        non_sequences=[rho, sigma],
+        n_steps=steps,
+        strict=True,
+        name="ma2",
+    )
+
+    def ref_logp(values, rho, sigma):
+        epsilon_tm2 = 0
+        epsilon_tm1 = 1
+        step_logps = np.zeros_like(values)
+        for t, value in enumerate(values):
+            mu = epsilon_tm1 * rho[0] + epsilon_tm2 * rho[1]
+            step_logps[t] = stats.norm.logpdf(value, mu, sigma)
+            epsilon_tm2 = epsilon_tm1
+            epsilon_tm1 = value - mu
+        return step_logps
+
+    ma2_vv = ma2.clone()
+    logp_expr = logp(ma2, ma2_vv)
+    assert_no_rvs(logp_expr)
+
+    ma2_test = rng.normal(size=(steps,))
+    rho_test = np.array([0.3, 0.7])
+    sigma_test = 0.9
+
+    np.testing.assert_array_almost_equal(
+        logp_expr.eval({ma2_vv: ma2_test, rho: rho_test, sigma: sigma_test}),
+        ref_logp(ma2_test, rho_test, sigma_test),
+    )