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):
@@ -552,6 +559,50 @@ def local_upcast_elemwise_constant_inputs(fgraph, node):
                 return rval
 
 
+@register_specialize
+@node_rewriter([Elemwise])
+def local_replace_broadcasted_constant(fgraph, node):
+    """Remove broadcasted constants from Elemwise graphs
+
+    Elemwise(scalar, ones((3, 4))) -> Alloc(Elemwise(scalar, ones((1, 1))), 3, 4)
+
+    This will avoid a 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
+
+
 @node_rewriter([Elemwise])
 def local_add_mul_fusion(fgraph, node):
     """Fuse consecutive add or mul in one such node with more inputs.

tests/tensor/rewriting/test_elemwise.py

@@ -178,6 +178,19 @@ def test_dimshuffle_lift_multi_out_elemwise(self):
         assert not local_dimshuffle_lift.transform(g, g.outputs[0].owner)
 
 
+def test_local_replace_broadcasted_constant():
+    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", "specialize"])
+    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")