Implement rewrite to inline Composite constants

Author: ricardoV94

pytensor/tensor/rewriting/elemwise.py

@@ -35,7 +35,7 @@
 from pytensor.tensor.math import exp
 from pytensor.tensor.rewriting.basic import 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_value
 
 
 class InplaceElemwiseOptimizer(GraphRewriter):
@@ -1203,6 +1203,44 @@ 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_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
+
+    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(
@@ -1243,6 +1281,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

@@ -1461,6 +1461,20 @@ def test_local_useless_composite_outputs():
     utt.assert_allclose(f([[np.nan]], [[1.0]], [[np.nan]]), [[0.0]])
 
 
+def test_local_inline_composite_constants():
+    x = vector("x")
+    out = exp(x + 2)
+
+    fn = pytensor.function([x], out, mode=get_default_mode().including("fusion"))
+
+    [node] = fn.maker.fgraph.apply_nodes
+    assert isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, Composite)
+    assert node.inputs == [x]
+
+    test_value = np.arange(5).astype(config.floatX)
+    np.testing.assert_allclose(fn(test_value), np.exp(test_value + 2))
+
+
 def test_local_useless_dimshuffle_makevector():
     a = scalar()
     x = MakeVector(config.floatX)(a)