Add a fusion rewrite for CAReduces with Elemwise inputs

Author: brandonwillard

aesara/scalar/basic.py

@@ -4048,7 +4048,7 @@ def __init__(self, inputs, outputs):
 
     @property
     def fn(self):
-        return self._fn
+        return None
 
     @property
     def inner_inputs(self):

aesara/tensor/rewriting/elemwise.py

@@ -11,11 +11,16 @@
 from aesara.graph.basic import Apply, Constant, io_toposort
 from aesara.graph.features import ReplaceValidate
 from aesara.graph.op import compute_test_value, get_test_value
-from aesara.graph.rewriting.basic import GraphRewriter, copy_stack_trace, node_rewriter
+from aesara.graph.rewriting.basic import (
+    GraphRewriter,
+    copy_stack_trace,
+    in2out,
+    node_rewriter,
+)
 from aesara.graph.rewriting.db import SequenceDB
 from aesara.graph.utils import InconsistencyError, MethodNotDefined, TestValueError
 from aesara.tensor.basic import MakeVector, alloc, cast, get_scalar_constant_value
-from aesara.tensor.elemwise import DimShuffle, Elemwise
+from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from aesara.tensor.exceptions import NotScalarConstantError
 from aesara.tensor.rewriting.basic import register_canonicalize, register_specialize
 from aesara.tensor.shape import shape_padleft
@@ -944,3 +949,58 @@ def local_useless_composite(fgraph, node):
         c = aes.Composite(inputs=comp.inputs, outputs=new_outputs)
         e = Elemwise(scalar_op=c)(*node.inputs, return_list=True)
         return dict(zip([node.outputs[i] for i in idx], e))
+
+
+@node_rewriter([CAReduce])
+def local_careduce_fusion(fgraph, node):
+    """Fuse a `CAReduce` applied to an `Elemwise`."""
+
+    (car_input,) = node.inputs
+    elm_node = car_input.owner
+
+    if elm_node is None or not isinstance(elm_node.op, Elemwise):
+        return False
+
+    elm_inputs = elm_node.inputs
+    elm_outputs = elm_node.outputs
+
+    if len(elm_inputs) > 1 or len(elm_outputs) > 1:
+        # TODO: Implement the multiple inputs case
+        return False
+
+    if len(fgraph.clients[elm_outputs[0]]) > 1:
+        return False
+
+    car_axis = node.op.axis
+    car_scalar_op = node.op.scalar_op
+    elm_scalar_op = elm_node.op.scalar_op
+
+    scalar_elm_inputs = [
+        aes.get_scalar_type(inp.type.dtype).make_variable() for inp in elm_inputs
+    ]
+    elm_output = elm_scalar_op(*scalar_elm_inputs)
+    # This input represents the previous value in the `CAReduce` binary reduction
+    carried_car_input = elm_output.type()
+    scalar_fused_outputs = [car_scalar_op(carried_car_input, elm_output)]
+
+    fused_scalar_op = aes.Composite(
+        inputs=[carried_car_input] + scalar_elm_inputs, outputs=scalar_fused_outputs
+    )
+
+    # The fused `Op` needs to look and behave like a `BinaryScalarOp`
+    # TODO: Generate a new `type` and make this relationship official?
+    fused_scalar_op.identity = car_scalar_op.identity
+    fused_scalar_op.nin = 2
+    fused_scalar_op.nout = 1
+
+    new_car_op = CAReduce(fused_scalar_op, car_axis)
+
+    return [new_car_op(*elm_inputs)]
+
+
+compile.optdb.register(  # type: ignore
+    "local_careduce_fusion",
+    in2out(local_careduce_fusion),
+    "fusion",
+    position=49,
+)

tests/tensor/rewriting/test_elemwise.py

@@ -1113,6 +1113,78 @@ def test_test_values(self, test_value):
                 f.maker.fgraph.outputs[0].tag.test_value, np.c_[[2.0]]
             )
 
+    @pytest.mark.parametrize("linker", ["cvm", "py"])
+    @pytest.mark.parametrize("axis", [None, 0, 1, (0, 1), (0, 1, 2)])
+    def test_CAReduce_single_input(self, linker, axis):
+        """Make sure that `CAReduce` and `Elemwise` fusions work with a single input."""
+
+        mode = Mode(linker=linker)
+        mode._optimizer = mode._optimizer.including(
+            "local_careduce_fusion",
+            "canonicalize",
+            "inplace",
+        )
+
+        x = tensor("floatX", shape=(None, None, None), name="x")
+        out = exp(x).sum(axis=axis)
+
+        out_fn = function([x], out, mode=mode)
+        (out_node,) = out_fn.maker.fgraph.toposort()
+
+        assert isinstance(getattr(out_node.op, "scalar_op"), aes.basic.Composite)
+
+        rng = np.random.default_rng(2320)
+        x_val = rng.random((4, 3, 2), dtype=config.floatX)
+
+        exp_res = np.exp(x_val).sum(axis=axis)
+
+        out_val = out_fn(x_val)
+        assert out_val.shape == exp_res.shape
+        assert np.allclose(out_val, exp_res)
+
+        # `Elemwise`s with more than one client shouldn't be rewritten
+        x = tensor("floatX", shape=(None, None, None), name="x")
+        exp_x = exp(x)
+        out = exp_x.sum(axis=axis) + exp(x)
+
+        out_fn = function([x], out, mode=mode)
+        out_nodes = out_fn.maker.fgraph.toposort()
+        assert not any(
+            isinstance(out_node.op.scalar_op, aes.basic.Composite)
+            for out_node in out_nodes
+            if hasattr(out_node.op, "scalar_op")
+        )
+
+    @pytest.mark.xfail(reason="Not implemented")
+    @pytest.mark.parametrize("linker", ["cvm", "py"])
+    @pytest.mark.parametrize("axis", [None, 0, 1, (0, 1), (0, 1, 2)])
+    def test_CAReduce_multiple_inputs(self, linker, axis):
+        """Make sure that `CAReduce` and `Elemwise` fusions work with multiple inputs."""
+
+        mode = Mode(linker=linker)
+        mode._optimizer = mode._optimizer.including(
+            "local_careduce_fusion",
+            "canonicalize",
+            "inplace",
+        )
+
+        x = tensor("floatX", shape=(None, None, None), name="x")
+        y = tensor("floatX", shape=(None, None, None), name="y")
+        out = (x + y).sum(axis=axis)
+
+        out_fn = function([x, y], out, mode=mode)
+        (out_node,) = out_fn.maker.fgraph.toposort()
+
+        assert isinstance(getattr(out_node.op, "scalar_op"), aes.basic.Composite)
+
+        rng = np.random.default_rng(2320)
+        x_val = rng.random((4, 3, 2), dtype=config.floatX)
+        y_val = rng.random((4, 3, 2), dtype=config.floatX)
+        exp_res = (x_val + y_val).sum(axis=axis)
+        out_val = out_fn(x_val, y_val)
+        assert out_val.shape == exp_res.shape
+        assert np.allclose(out_val, exp_res)
+
 
 class TimesN(aes.basic.UnaryScalarOp):
     """