Add rewrite for Sum(MakeVector)

Author: aseyboldt

pytensor/tensor/rewriting/basic.py

@@ -43,6 +43,7 @@
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.extra_ops import broadcast_shape, broadcast_to
+from pytensor.tensor.math import Sum, add
 from pytensor.tensor.math import all as at_all
 from pytensor.tensor.math import eq
 from pytensor.tensor.shape import Shape_i
@@ -956,6 +957,30 @@ def local_join_make_vector(fgraph, node):
         return [ret]
 
 
+@register_specialize
+@register_canonicalize
+@register_useless
+@node_rewriter([Sum])
+def local_sum_make_vector(fgraph, node):
+    """A sum of a MakeVector node is just the sum of the elements."""
+    (array,) = node.inputs
+
+    if array.owner is None:
+        return
+
+    if not isinstance(array.owner.op, MakeVector):
+        return
+
+    if node.op.axis not in [None, 0, -1]:
+        return
+
+    elements = array.owner.inputs
+    dtype = node.op.acc_dtype
+    element_sum = add(*[cast(value, dtype) for value in elements])
+
+    return [as_tensor_variable(element_sum)]
+
+
 @register_useless("local_remove_switch_const_cond")
 @register_canonicalize("fast_compile", "local_remove_switch_const_cond")
 @register_specialize

tests/tensor/rewriting/test_basic.py

@@ -19,6 +19,7 @@
 from pytensor.graph.rewriting.utils import rewrite_graph
 from pytensor.printing import debugprint, pprint
 from pytensor.raise_op import Assert, CheckAndRaise
+from pytensor.scalar.basic import Add
 from pytensor.tensor.basic import (
     Alloc,
     Join,
@@ -102,6 +103,7 @@
     values_eq_approx_remove_nan,
     vector,
 )
+from pytensor.tensor.var import TensorVariable
 from tests import unittest_tools as utt
 
 
@@ -1300,6 +1302,20 @@ def test_local_join_make_vector():
     assert check_stack_trace(f, ops_to_check="all")
 
 
+def test_local_sum_make_vector():
+    a, b, c = scalars("abc")
+    mv = MakeVector(config.floatX)
+    output = mv(a, b, c).sum()
+
+    func = function([a, b, c], output)
+
+    elemwise = func.maker.fgraph.outputs[0].owner
+    # The MakeVector op should be optimized away, so we just
+    # take the sum of the scalars.
+    assert elemwise.inputs[0].name == "a"
+    assert isinstance(elemwise.inputs[0], TensorVariable)
+
+
 @pytest.mark.parametrize(
     "dtype",
     [