Introduce RandomVariable push-out optimization

Author: brandonwillard

symbolic_pymc/theano/opt.py

@@ -5,7 +5,9 @@
 
 from functools import wraps
 
-from theano.gof.opt import LocalOptimizer
+from theano.gof.opt import LocalOptimizer, local_optimizer
+from theano.scan_module.scan_op import Scan
+from theano.scan_module.scan_utils import scan_args as ScanArgs
 
 from unification import var, variables
 
@@ -14,6 +16,7 @@
 from etuples.core import ExpressionTuple
 
 from .meta import MetaSymbol
+from .ops import RandomVariable
 
 
 def eval_and_reify_meta(x):
@@ -219,3 +222,54 @@ def transform(self, node):
             return new_node
         else:
             return False
+
+
+@local_optimizer([Scan])
+def push_out_rvs_from_scan(node):
+    """Push `RandomVariable`s out of `Scan` nodes.
+
+    When `RandomVariable`s are created within the inner-graph of a `Scan` and
+    are not output to the outer-graph, we "push" them out of the inner-graph.
+    This helps us produce an outer-graph in which all the relevant `RandomVariable`s
+    are accessible (e.g. for constructing a log-likelihood graph).
+    """
+    if not isinstance(node.op, Scan):
+        return False
+
+    scan_args = ScanArgs(node.inputs, node.outputs, node.op.inputs, node.op.outputs, node.op.info)
+
+    # Find the un-output `RandomVariable`s created in the inner-graph
+    clients = {}
+    local_fgraph_topo = theano.gof.graph.io_toposort(
+        scan_args.inner_inputs, scan_args.inner_outputs, clients=clients
+    )
+    unpushed_inner_rvs = []
+    for n in local_fgraph_topo:
+        if isinstance(n.op, RandomVariable):
+            unpushed_inner_rvs.extend([c for c in clients[n] if c not in scan_args.inner_outputs])
+
+    if len(unpushed_inner_rvs) == 0:
+        return False
+
+    # Add the new outputs to the inner and outer graphs
+    scan_args.inner_out_nit_sot.extend(unpushed_inner_rvs)
+
+    if len(scan_args.outer_in_nit_sot) == 0:
+        raise ValueError("No outer-graph inputs nit-sots!")
+
+    # Just like `theano.scan`, we simply copy/repeat the existing nit-sot
+    # outer-graph input value, which represents the actual size of the output
+    # tensors.  Apparently, the value needs to be duplicated for all nit-sots.
+    # FYI: This is what increments the nit-sot values in `scan_args.info`, as
+    # well.
+    # TODO: Can we just use `scan_args.n_steps`?
+    scan_args.outer_in_nit_sot.extend(scan_args.outer_in_nit_sot[0:1] * len(unpushed_inner_rvs))
+
+    op = Scan(scan_args.inner_inputs, scan_args.inner_outputs, scan_args.info)
+    outputs = list(op(*scan_args.outer_inputs))
+
+    # Return only the replacements for the original `node.outputs`
+    new_inner_out_idx = [scan_args.inner_outputs.index(i) for i in unpushed_inner_rvs]
+    _ = [outputs.pop(op.var_mappings["outer_out_from_inner_out"][i]) for i in new_inner_out_idx]
+
+    return dict(zip(node.outputs, outputs))

tests/theano/test_opt.py

@@ -1,3 +1,5 @@
+import numpy as np
+import theano
 import theano.tensor as tt
 
 from unification import var
@@ -9,10 +11,16 @@
 
 from theano.gof.opt import EquilibriumOptimizer
 from theano.gof.graph import inputs as tt_inputs
+from theano.scan_module.scan_op import Scan
 
 from symbolic_pymc.theano.meta import mt
-from symbolic_pymc.theano.opt import KanrenRelationSub, FunctionGraph
+from symbolic_pymc.theano.opt import (
+    KanrenRelationSub,
+    FunctionGraph,
+    push_out_rvs_from_scan,
+)
 from symbolic_pymc.theano.utils import optimize_graph
+from symbolic_pymc.theano.random_variables import CategoricalRV, DirichletRV, NormalRV
 
 
 def test_kanren_opt():
@@ -58,3 +66,64 @@ def distributes(in_lv, out_lv):
     assert fgraph_opt.owner.inputs[1].owner.op == tt.add
     assert isinstance(fgraph_opt.owner.inputs[1].owner.inputs[0].owner.op, tt.Dot)
     assert isinstance(fgraph_opt.owner.inputs[1].owner.inputs[1].owner.op, tt.Dot)
+
+
+def test_push_out_rvs():
+    theano.config.cxx = ""
+    theano.config.mode = "FAST_COMPILE"
+    tt.config.compute_test_value = "warn"
+
+    rng_state = np.random.RandomState(np.random.MT19937(np.random.SeedSequence(1234)))
+    rng_tt = theano.shared(rng_state, name="rng", borrow=True)
+    rng_tt.tag.is_rng = True
+    rng_tt.default_update = rng_tt
+
+    N_tt = tt.iscalar("N")
+    N_tt.tag.test_value = 10
+    M_tt = tt.iscalar("M")
+    M_tt.tag.test_value = 2
+
+    mus_tt = tt.matrix("mus_t")
+    mus_tt.tag.test_value = np.stack([np.arange(0.0, 10), np.arange(0.0, -10, -1)], axis=-1).astype(
+        theano.config.floatX
+    )
+
+    sigmas_tt = tt.ones((N_tt,))
+    Gamma_rv = DirichletRV(tt.ones((M_tt, M_tt)), rng=rng_tt, name="Gamma")
+
+    # In this case, `Y_t` depends on `S_t` and `S_t` is not output.  Our
+    # push-out optimization should create a new `Scan` that also outputs each
+    # `S_t`.
+    def scan_fn(mus_t, sigma_t, Gamma_t, rng):
+        S_t = CategoricalRV(Gamma_t[0], rng=rng, name="S_t")
+        Y_t = NormalRV(mus_t[S_t], sigma_t, rng=rng, name="Y_t")
+        return Y_t
+
+    Y_rv, _ = theano.scan(
+        fn=scan_fn,
+        sequences=[mus_tt, sigmas_tt],
+        non_sequences=[Gamma_rv, rng_tt],
+        outputs_info=[{}],
+        strict=True,
+        name="scan_rv",
+    )
+    Y_rv.name = "Y_rv"
+
+    orig_scan_op = Y_rv.owner.op
+    assert len(Y_rv.owner.outputs) == 2
+    assert isinstance(orig_scan_op, Scan)
+    assert len(orig_scan_op.outputs) == 2
+    assert orig_scan_op.outputs[0].owner.op == NormalRV
+    assert isinstance(orig_scan_op.outputs[1].type, tt.raw_random.RandomStateType)
+
+    fgraph = FunctionGraph(tt_inputs([Y_rv]), [Y_rv], clone=True)
+    pushoutrvs_opt = EquilibriumOptimizer([push_out_rvs_from_scan], max_use_ratio=10)
+    fgraph_opt = optimize_graph(fgraph, pushoutrvs_opt, return_graph=True)
+
+    # There should now be a new output for all the `S_t`
+    new_scan = fgraph_opt.outputs[0].owner
+    assert len(new_scan.outputs) == 3
+    assert isinstance(new_scan.op, Scan)
+    assert new_scan.op.outputs[0].owner.op == NormalRV
+    assert new_scan.op.outputs[1].owner.op == CategoricalRV
+    assert isinstance(new_scan.op.outputs[2].type, tt.raw_random.RandomStateType)