Implement rewrite to inline Composite constants

Author: ricardoV94

pytensor/tensor/rewriting/elemwise.py

@@ -31,7 +31,7 @@
     alloc,
     cast,
     constant,
-    get_underlying_scalar_constant_value,
+    get_underlying_scalar_constant_value, constant,
 )
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
 from pytensor.tensor.exceptions import NotScalarConstantError
@@ -1254,6 +1254,49 @@ def local_careduce_fusion(fgraph, node):
     return [new_car_op(*elm_inputs)]
 
 
+@register_specialize
+@node_rewriter([Elemwise])
+def local_inline_composite_constants(fgraph, node):
+    """Inline scalar constants in Composite graphs."""
+    composite_op = node.op.scalar_op
+
+    if not isinstance(composite_op, aes.Composite):
+        return None
+
+    new_outer_inputs = []
+    new_inner_inputs = []
+    inner_replacements = {}
+    for outer_inp, inner_inp in zip(node.inputs, composite_op.fgraph.inputs):
+        # Complex variables don't have a `c_literal` that can be inlined
+        if "complex" not in outer_inp.type.dtype:
+            unique_value = get_unique_constant_value(outer_inp)
+            if unique_value is not None:
+                inner_replacements[inner_inp] = aes.constant(
+                    unique_value, dtype=inner_inp.dtype
+                )
+                continue
+        new_outer_inputs.append(outer_inp)
+        new_inner_inputs.append(inner_inp)
+
+    if not inner_replacements:
+        return None
+
+    new_inner_outs = clone_replace(
+        composite_op.fgraph.outputs, replace=inner_replacements
+    )
+    new_composite_op = aes.Composite(new_inner_inputs, new_inner_outs)
+    new_outputs = Elemwise(new_composite_op).make_node(*new_outer_inputs).outputs
+
+    # Some of the inlined constants were broadcasting the output shape
+    # This should be taken care of by `local_replace_broadcasted_constant`
+    for old_out, new_out in zip(node.outputs, new_outputs):
+        if old_out.type.broadcastable != new_out.type.broadcastable:
+            return None
+
+    copy_stack_trace(node.outputs, new_outputs)
+    return new_outputs
+
+
 # Register fusion database just before AddDestroyHandler(49.5) (inplace rewrites)
 fuse_seqopt = SequenceDB()
 compile.optdb.register(
@@ -1294,6 +1337,13 @@ def local_careduce_fusion(fgraph, node):
     "fusion",
     position=10,
 )
+fuse_seqopt.register(
+    "local_inline_composite_constants",
+    in2out(local_inline_composite_constants),
+    "fast_run",
+    "fusion",
+    position=20,
+)
 
 
 def _rebuild_partial_2f1grad_loop(node, wrt):

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, equal_computations
+from pytensor.graph.basic import Constant, ancestors, 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
@@ -1474,6 +1474,25 @@ def test_local_useless_composite_outputs():
     utt.assert_allclose(f([[np.nan]], [[1.0]], [[np.nan]]), [[0.0]])
 
 
+@pytest.mark.parametrize("const_shape", [(), (1,), (5,), (1, 5), (2, 5)])
+@pytest.mark.parametrize("op, np_op", [(at.pow, np.power), (at.add, np.add)])
+def test_local_inline_composite_constants(op, np_op, const_shape):
+    x = vector("x")
+    const = np.full(shape=const_shape, fill_value=2.5)
+    out = exp(op(x, const))
+
+    fn = pytensor.function([x], out, mode=get_default_mode().including("specialize", "fusion"))
+    [node] = [
+        node for node in fn.maker.fgraph.apply_nodes if isinstance(node.op, Elemwise)
+    ]
+    assert isinstance(node.op.scalar_op, Composite)
+    assert len(node.inputs) == 1
+    assert x in ancestors(node.inputs)
+
+    test_value = np.arange(5).astype(config.floatX)
+    np.testing.assert_allclose(fn(test_value), np.exp(np_op(test_value, const)))
+
+
 def test_local_useless_dimshuffle_makevector():
     a = scalar()
     x = MakeVector(config.floatX)(a)