Add rewrite to remove broadcasted constants from Elemwise graphs

Author: ricardoV94

pytensor/tensor/rewriting/elemwise.py

@@ -4,6 +4,8 @@
 from typing import DefaultDict, Generator, List, Set, Tuple, TypeVar
 from warnings import warn
 
+import numpy as np
+
 import pytensor
 import pytensor.scalar.basic as aes
 from pytensor import clone_replace, compile
@@ -28,14 +30,19 @@
     MakeVector,
     alloc,
     cast,
+    constant,
     get_underlying_scalar_constant_value,
 )
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.math import exp
-from pytensor.tensor.rewriting.basic import register_canonicalize, register_specialize
+from pytensor.tensor.rewriting.basic import (
+    broadcast_like,
+    register_canonicalize,
+    register_specialize,
+)
 from pytensor.tensor.shape import shape_padleft
-from pytensor.tensor.var import TensorConstant
+from pytensor.tensor.var import TensorConstant, get_unique_constant_value
 
 
 class InplaceElemwiseOptimizer(GraphRewriter):
@@ -1296,6 +1303,65 @@ def local_inline_composite_constants(fgraph, node):
 )
 
 
+@node_rewriter([Elemwise])
+def local_replace_broadcasted_constants(fgraph, node):
+    """Remove broadcasted constants from Elemwise graphs
+
+    Elemwise(matrix, ones((3, 4))) -> Elemwise(vector, ones((1, 1)))
+
+    In cases where the constant influenced the final shape of the Elemwise operation
+    We broadcast (via alloc) the new Elemwise result:
+
+    Elemwise(row, ones((3, 4))) -> Alloc(Elemwise(row, ones((1, 1))), 3, 4)
+
+    This will avoid useless iterations over constant arrays.
+    """
+    if len(node.inputs) == 1:
+        return None
+
+    new_elem_inps = []
+    ndims = node.outputs[0].type.ndim
+    found_const = False
+    for inp in node.inputs:
+        # If input has non-broadcastable dims
+        if not all(b for b in inp.type.broadcastable):
+            constant_value = get_unique_constant_value(inp)
+            if constant_value is not None:
+                constant_value = np.expand_dims(
+                    constant_value, axis=tuple(range(ndims))
+                ).astype(inp.type.dtype)
+                new_elem_inps.append(constant(constant_value))
+                found_const = True
+                continue
+
+        new_elem_inps.append(inp)
+
+    if not found_const:
+        return None
+
+    new_outs = node.op.make_node(*new_elem_inps).outputs
+
+    # The constants were needed to enforce the output shape
+    if node.outputs[0].type.broadcastable != new_outs[0].type.broadcastable:
+        new_outs = [
+            broadcast_like(new_out, template=node.outputs[0], fgraph=fgraph)
+            for new_out in new_outs
+        ]
+
+    copy_stack_trace(node.outputs, new_outs)
+    return new_outs
+
+
+# We register this immediately after the fusion database.
+# We don't want Allocs to break up the fusion rewrites
+compile.optdb.register(
+    "local_replace_broadcasted_constants",
+    in2out(local_replace_broadcasted_constants),
+    "fast_run",
+    position=49.01,
+)
+
+
 def _rebuild_partial_2f1grad_loop(node, wrt):
     a, b, c, log_z, sign_z = node.inputs[-5:]
     z = exp(log_z) * sign_z

tests/tensor/rewriting/test_elemwise.py

@@ -12,7 +12,7 @@
 from pytensor.compile.mode import Mode, get_default_mode
 from pytensor.configdefaults import config
 from pytensor.gradient import grad
-from pytensor.graph.basic import Constant, ancestors, equal_computations
+from pytensor.graph.basic import Constant, equal_computations
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.rewriting.basic import check_stack_trace, out2in
 from pytensor.graph.rewriting.db import RewriteDatabaseQuery
@@ -178,6 +178,21 @@ def test_dimshuffle_lift_multi_out_elemwise(self):
         assert not local_dimshuffle_lift.transform(g, g.outputs[0].owner)
 
 
+def test_local_replace_broadcasted_constants():
+    const = np.full(shape=(2, 5), fill_value=2.6)
+    x = scalar("x")
+    out = at.power(x, const)
+    new_out = rewrite_graph(
+        out, include=["ShapeOpt", "local_replace_broadcasted_constants"]
+    )
+    ref_out = at.alloc(
+        at.power(x, [[2.6]]),
+        at.constant(2, dtype="int64"),
+        at.constant(5, dtype="int64"),
+    )
+    assert equal_computations([new_out], [ref_out])
+
+
 def test_local_useless_dimshuffle_in_reshape():
     vec = TensorType(dtype="float64", shape=(None,))("vector")
     mat = TensorType(dtype="float64", shape=(None, None))("mat")