Skip to content
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.

Fuse CAReduces and Elemwises #1116

Description

@brandonwillard

We aren't currently converting CAReduces of Elemwises into single CAReduces:

import aesara
import aesara.tensor as at


x = at.vector("x")
y = at.exp(x).sum()

y_fn = aesara.function([x], y)

aesara.dprint(y_fn)
# Sum{acc_dtype=float64} [id A] 1
#  |Elemwise{exp,no_inplace} [id B] 0
#    |x [id C]

As the example shows, we first compute the Elemwise then immediately reduce it. Instead, we could fuse the CAReduce and Elemwise scalar Ops (i.e. create a Composite scalar Op corresponding to add(exp(x), exp(y))) and use that in a single CAReduce node that avoids the need for intermediate storage.

The rewrites required to do this would incorporate most of the same logic from aesara.tensor.basic_opt.FusionOptimizer, but, instead of producing another Elemwise, it would result in a new CAReduce.

Metadata

Metadata

Assignees

No one assigned

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions