Skip to content

Optimize reductions of broadcasting (alloc) #378

Open
@ricardoV94

Description

@ricardoV94

Description

import pytensor
import pytensor.tensor as pt

x = pt.vector("x")
out = pt.sum(pt.ones((5, 3)) * x)
fn = pytensor.function([x], out)
pytensor.dprint(fn, print_type=True)
Sum{axes=None} [id A] <Scalar(float64, shape=())> 1
 └─ Alloc [id B] <Matrix(float64, shape=(5, 3))> 0
    ├─ x [id C] <Vector(float64, shape=(?,))>
    ├─ 5 [id D] <Scalar(int64, shape=())>
    └─ 3 [id E] <Scalar(int64, shape=())>

Could instead be x * 15. Boolean reductions (all, any) are even simpler.

Maybe something clever can be done with axis, but as a start we can cover just the cases with axis=None

Conceptually related to #59

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