This repository was archived by the owner on Nov 17, 2025. It is now read-only.
This repository was archived by the owner on Nov 17, 2025. It is now read-only.
Allow merging graphs with multiple clients in FusionOptimizer #1237
Description
Right now, Elemwise composite graphs will be split at nodes that have multiple clients. This can be avoided by allowing the fusion rewriter to replace multiple outputs. One example where this is very common is in value_and_grad functions.
import aesara
import aesara.tensor as at
x = at.vector("x")
y = at.exp(x/5)
w = y + 1
z = y * 2
f = aesara.function([x], [z, w])
aesara.dprint(f)Elemwise{Mul}[(0, 1)] [id A] 2
|TensorConstant{(1,) of 2.0} [id B]
|Elemwise{Composite{exp((i0 * i1))}} [id C] 0
|TensorConstant{(1,) of 0.2} [id D]
|x [id E]
Elemwise{add,no_inplace} [id F] 1
|TensorConstant{(1,) of 1.0} [id G]
|Elemwise{Composite{exp((i0 * i1))}} [id C] 0
f = aesara.function([x], [z])
aesara.dprint(f)Elemwise{Composite{(i0 * exp((i1 * i2)))}} [id A] 0
|TensorConstant{(1,) of 2.0} [id B]
|TensorConstant{(1,) of 0.2} [id C]
|x [id D]
Discussed with @aseyboldt and @aesara-devs/aesara
See rationale and proposed changes in #1237 (comment)