Inline SymbolicOps at the start of the fusion pass so their bodies can fuse

pytensor/compile/mode.py

@@ -524,7 +524,9 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 )
 PYTORCH = Mode(
     PytorchLinker(),
-    RewriteDatabaseQuery(include=["fast_run"]),
+    # `Composite` has no PyTorch dispatch (it has no `nfunc_spec` to look up in torch),
+    # and torch fuses its own kernels anyway.
+    RewriteDatabaseQuery(include=["fast_run"], exclude=["fusion"]),
 )
 
 MLX = Mode(

pytensor/tensor/rewriting/ofg.py

@@ -1,8 +1,17 @@
 from pytensor.compile.rewriting import inline_ofg_node
 from pytensor.graph import node_rewriter
+from pytensor.graph.rewriting.basic import MergeOptimizer, dfs_rewriter
 from pytensor.tensor.basic import AllocDiag
 from pytensor.tensor.rewriting.basic import register_specialize
-from pytensor.tensor.special import LogAddExp, LogSumExp, XLog1PY, XLogY
+from pytensor.tensor.rewriting.elemwise import fuse_seqopt
+from pytensor.tensor.special import (
+    LogAddExp,
+    LogSoftmax,
+    LogSumExp,
+    Softmax,
+    XLog1PY,
+    XLogY,
+)
 
 
 @register_specialize("inline_ofg")
@@ -32,3 +41,43 @@ def late_inline_OpFromGraph(fgraph, node):
 
     """
     return inline_ofg_node(node)
+
+
+@node_rewriter([Softmax, LogSoftmax, LogSumExp, LogAddExp, XLogY, XLog1PY])
+def inline_symbolic_for_fusion(fgraph, node):
+    """Inline `SymbolicOp`s so their bodies can fuse.
+
+    The inner graph is otherwise compiled on its own, making the op a fusion barrier
+    both inside and across its boundary. Ops are listed explicitly, as some are worth
+    keeping opaque anyway (`KroneckerProduct` is: its body is one broadcasting `mul`
+    behind a `reshape`). Backends that dispatch these ops directly (JAX, PyTorch and MLX
+    have their own `Softmax`) exclude the whole fusion pass, so they keep them whole.
+
+    Ops that `late_inline_OpFromGraph` already inlines at specialize are listed here
+    too, since fusion also runs in modes that skip specialize.
+    """
+    return inline_ofg_node(node)
+
+
+# Fusion is the only thing that needs the body, so inline as its first step: every
+# Op-level rewrite still matches on the op.
+fuse_seqopt.register(
+    "inline_symbolic_for_fusion",
+    dfs_rewriter(inline_symbolic_for_fusion),
+    "fast_run",
+    "fusion",
+    position=0,
+)
+# Each op is inlined from its own inner graph, so two ops over the same inputs (a
+# `Softmax` and a `LogSoftmax` over the same logits) yield structurally identical but
+# distinct copies of everything they share. The `merge2` pass cannot clean this up: it
+# sits at the same optdb position as the fusion pass, where ties are broken by name, so
+# it only runs once fusion has pulled both copies into the `Composite`.
+fuse_seqopt.register(
+    "merge_inlined_symbolic",
+    MergeOptimizer(),
+    "fast_run",
+    "fusion",
+    "merge",
+    position=0.5,
+)

tests/tensor/rewriting/test_elemwise.py

@@ -1458,6 +1458,60 @@ def test_many_fused_subgraphs(self):
         fgraph = FunctionGraph([x_row, x_mat], [m])
         FusionOptimizer().apply(fgraph)
 
+    def test_fuse_across_symbolic_op(self):
+        # A SymbolicOp with no backend dispatch is a fusion barrier: its body is
+        # compiled on its own, so the ops around it cannot join the ops inside it.
+        # Inlining it at the start of the fusion pass lets the whole chain fuse --
+        # the subtraction from the log_softmax body ends up in the same Composite
+        # as the surrounding exp and add.
+        rng = np.random.default_rng(75)
+        x = pt.tensor("x", shape=(5, 4))
+        out = exp(pt.special.log_softmax(x * 2.0, axis=-1)) + 1.0
+
+        def output_scalar_ops(fn):
+            scalar_op = fn.maker.fgraph.outputs[0].owner.op.scalar_op
+            assert isinstance(scalar_op, Composite)
+            return {node.op for node in scalar_op.fgraph.apply_nodes}
+
+        fn = function([x], out, mode=self.mode)
+        assert output_scalar_ops(fn) == {ps.add, ps.exp, ps.sub}
+
+        # Without the inlining the body stays sealed off, so only the outer ops fuse
+        barrier_fn = function(
+            [x], out, mode=self.mode.excluding("inline_symbolic_for_fusion")
+        )
+        assert output_scalar_ops(barrier_fn) == {ps.add, ps.exp}
+
+        x_val = rng.normal(size=(5, 4)).astype(config.floatX)
+        np.testing.assert_allclose(fn(x_val), barrier_fn(x_val), rtol=1e-6)
+
+    def test_merge_inlined_symbolic_ops(self):
+        # Each SymbolicOp is inlined from its own inner graph, so a Softmax and a
+        # LogSoftmax over the same logits come out with distinct copies of everything
+        # their bodies share. Merging them before fusion is what keeps the Composite
+        # from evaluating the exponential over each copy.
+        x = pt.tensor("x", shape=(5, 4))
+        outs = [pt.special.softmax(x, axis=-1), pt.special.log_softmax(x, axis=-1)]
+
+        def n_exp(fn):
+            scalar_ops = []
+            for node in fn.maker.fgraph.toposort():
+                scalar_op = getattr(node.op, "scalar_op", None)
+                if isinstance(scalar_op, Composite):
+                    scalar_ops.extend(n.op for n in scalar_op.fgraph.apply_nodes)
+                elif scalar_op is not None:
+                    scalar_ops.append(scalar_op)
+            return scalar_ops.count(ps.exp)
+
+        fn = function([x], outs, mode=self.mode)
+        assert n_exp(fn) == 1
+
+        # Without the merge the two copies survive into the Composite
+        unmerged_fn = function(
+            [x], outs, mode=self.mode.excluding("merge_inlined_symbolic")
+        )
+        assert n_exp(unmerged_fn) == 2
+
 
 class TimesN(ps.basic.UnaryScalarOp):
     """

tests/tensor/rewriting/test_special.py

@@ -26,7 +26,10 @@
 from tests.unittest_tools import RewriteTester
 
 
-_fast_run_rewrites = RewriteDatabaseQuery(include=["fast_run"])
+# Fusion inlines `Softmax` and `LogSoftmax`, which would hide the ops these tests check for
+_fast_run_rewrites = RewriteDatabaseQuery(
+    include=["fast_run"], exclude=["inline_symbolic_for_fusion"]
+)
 _fast_run_rewrites = optdb.query(_fast_run_rewrites)
 
 

tests/tensor/test_special.py

@@ -69,7 +69,11 @@ def test_infer_shape(self):
         rng = np.random.default_rng(utt.fetch_seed())
         admat_val = rng.random((3, 4)).astype(config.floatX)
         self._compile_and_check(
-            [admat], [Softmax(axis=-1)(admat)], [admat_val], Softmax
+            [admat],
+            [Softmax(axis=-1)(admat)],
+            [admat_val],
+            Softmax,
+            excluding=["inline_symbolic_for_fusion"],
         )
 
     def test_vector_perform(self):