Introduce graph rewrite for mixture sub-graphs defined via Switch Op

Author: larryshamalama

aeppl/mixture.py

@@ -11,14 +11,17 @@
     pre_greedy_node_rewriter,
 )
 from aesara.ifelse import ifelse
+from aesara.scalar.basic import Switch
 from aesara.tensor.basic import Join, MakeVector
+from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.random.rewriting import (
     local_dimshuffle_rv_lift,
     local_subtensor_rv_lift,
 )
 from aesara.tensor.shape import shape_tuple
 from aesara.tensor.subtensor import (
     as_index_literal,
+    as_nontensor_scalar,
     get_canonical_form_slice,
     is_basic_idx,
 )
@@ -251,7 +254,6 @@ def mixture_replace(fgraph, node):
     From these terms, new terms ``Z_rv[i] = mixture_comps[i][i == I_rv]`` are
     created for each ``i`` in ``enumerate(mixture_comps)``.
     """
-
     rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
 
     if rv_map_feature is None:
@@ -303,6 +305,56 @@ def mixture_replace(fgraph, node):
     return [new_mixture_rv]
 
 
+@node_rewriter((Elemwise,))
+def switch_mixture_replace(fgraph, node):
+    rv_map_feature = getattr(fgraph, "preserve_rv_mappings", None)
+
+    if rv_map_feature is None:
+        return None  # pragma: no cover
+
+    if not isinstance(node.op.scalar_op, Switch):
+        return None  # pragma: no cover
+
+    old_mixture_rv = node.default_output()
+    # idx, component_1, component_2 = node.inputs
+
+    mixture_rvs = []
+
+    for component_rv in node.inputs[1:]:
+        if not (
+            component_rv.owner
+            and isinstance(component_rv.owner.op, MeasurableVariable)
+            and component_rv not in rv_map_feature.rv_values
+        ):
+            return None
+        new_node = assign_custom_measurable_outputs(component_rv.owner)
+        out_idx = component_rv.owner.outputs.index(component_rv)
+        new_comp_rv = new_node.outputs[out_idx]
+        mixture_rvs.append(new_comp_rv)
+
+    mix_op = MixtureRV(
+        2,
+        old_mixture_rv.dtype,
+        old_mixture_rv.broadcastable,
+    )
+    new_node = mix_op.make_node(
+        *([NoneConst, as_nontensor_scalar(node.inputs[0])] + mixture_rvs)
+    )
+
+    new_mixture_rv = new_node.default_output()
+
+    if aesara.config.compute_test_value != "off":
+        if not hasattr(old_mixture_rv.tag, "test_value"):
+            compute_test_value(node)
+
+        new_mixture_rv.tag.test_value = old_mixture_rv.tag.test_value
+
+    if old_mixture_rv.name:
+        new_mixture_rv.name = f"{old_mixture_rv.name}-mixture"
+
+    return [new_mixture_rv]
+
+
 @_logprob.register(MixtureRV)
 def logprob_MixtureRV(
     op, values, *inputs: Optional[Union[TensorVariable, slice]], name=None, **kwargs
@@ -368,7 +420,8 @@ def logprob_MixtureRV(
 logprob_rewrites_db.register(
     "mixture_replace",
     EquilibriumGraphRewriter(
-        [mixture_replace], max_use_ratio=aesara.config.optdb__max_use_ratio
+        [mixture_replace, switch_mixture_replace],
+        max_use_ratio=aesara.config.optdb__max_use_ratio,
     ),
     0,
     "basic",

tests/test_mixture.py

@@ -3,13 +3,14 @@
 import numpy as np
 import pytest
 import scipy.stats.distributions as sp
-from aesara.graph.basic import Variable
+from aesara.graph.basic import Variable, equal_computations
 from aesara.tensor.random.basic import CategoricalRV
 from aesara.tensor.shape import shape_tuple
 from aesara.tensor.subtensor import as_index_constant
 
 from aeppl.joint_logprob import factorized_joint_logprob, joint_logprob
-from aeppl.mixture import expand_indices
+from aeppl.mixture import MixtureRV, expand_indices
+from aeppl.rewriting import construct_ir_fgraph
 from tests.test_logprob import scipy_logprob
 from tests.utils import assert_no_rvs
 
@@ -67,7 +68,14 @@ def create_mix_model(size, axis):
 
 
 @aesara.config.change_flags(compute_test_value="warn")
-def test_compute_test_value():
+@pytest.mark.parametrize(
+    "op_constructor",
+    [
+        lambda _I, _X, _Y: at.stack([_X, _Y])[_I],
+        lambda _I, _X, _Y: at.switch(_I, _X, _Y),
+    ],
+)
+def test_compute_test_value(op_constructor):
 
     srng = at.random.RandomStream(29833)
 
@@ -82,7 +90,7 @@ def test_compute_test_value():
     i_vv = I_rv.clone()
     i_vv.name = "i"
 
-    M_rv = at.stack([X_rv, Y_rv])[I_rv]
+    M_rv = op_constructor(I_rv, X_rv, Y_rv)
     M_rv.name = "M"
 
     m_vv = M_rv.clone()
@@ -705,3 +713,49 @@ def test_mixture_with_DiracDelta():
     logp_res = factorized_joint_logprob({M_rv: m_vv, I_rv: i_vv})
 
     assert m_vv in logp_res
+
+
+def test_switch_mixture():
+    srng = at.random.RandomStream(29833)
+
+    X_rv = srng.normal(-10.0, 0.1, name="X")
+    Y_rv = srng.normal(10.0, 0.1, name="Y")
+
+    I_rv = srng.bernoulli(0.5, name="I")
+    i_vv = I_rv.clone()
+    i_vv.name = "i"
+
+    Z1_rv = at.switch(I_rv, X_rv, Y_rv)
+    z_vv = Z1_rv.clone()
+    z_vv.name = "z1"
+
+    fgraph, _, _ = construct_ir_fgraph({Z1_rv: z_vv, I_rv: i_vv})
+
+    assert isinstance(fgraph.outputs[0].owner.op, MixtureRV)
+    assert not hasattr(
+        fgraph.outputs[0].tag, "test_value"
+    )  # aesara.config.compute_test_value == "off"
+    assert fgraph.outputs[0].name is None
+
+    Z1_rv.name = "Z1"
+
+    fgraph, _, _ = construct_ir_fgraph({Z1_rv: z_vv, I_rv: i_vv})
+
+    assert fgraph.outputs[0].name == "Z1-mixture"
+
+    # building the identical graph but with a stack to check that mixture computations are identical
+
+    Z2_rv = at.stack((X_rv, Y_rv))[I_rv]
+
+    fgraph2, _, _ = construct_ir_fgraph({Z2_rv: z_vv, I_rv: i_vv})
+
+    assert equal_computations(fgraph.outputs, fgraph2.outputs)
+
+    z1_logp = joint_logprob({Z1_rv: z_vv, I_rv: i_vv})
+    z2_logp = joint_logprob({Z2_rv: z_vv, I_rv: i_vv})
+
+    # below should follow immediately from the equal_computations assertion above
+    assert equal_computations([z1_logp], [z2_logp])
+
+    np.testing.assert_almost_equal(0.69049938, z1_logp.eval({z_vv: -10, i_vv: 0}))
+    np.testing.assert_almost_equal(0.69049938, z2_logp.eval({z_vv: -10, i_vv: 0}))