Skip to content

Consider maknig more use of python match in grah rewrites #1455

Open
@ricardoV94

Description

@ricardoV94

Description

This may make code a bit cleaner than our usual checks like x.owner is not None and isinstance(x.owner.op, Elemwise) and isinstance(x.owner.op.scalar_op, ps.Add)

from pytensor.graph.basic import Apply
from pytensor.tensor.elemwise import Elemwise
import pytensor.scalar as ps
import pytensor.tensor as pt

def match_owner(x):
    match x.owner:
        case Apply(op=Elemwise(scalar_op=ps.Add())):
            return True
    return False

x = pt.vector("x")
y = x + x
z = pt.exp(y)
print(f"{match_owner(x)=}")
print(f"{match_owner(y)=}")
print(f"{match_owner(z)=}")
# match_owner(x)=False
# match_owner(y)=True
# match_owner(z)=False

Even more for Ops that can be finely parametrized like Solve(b_ndim=1) in cases that distinction matters.

OTOH checking if isinstance(x.owner, Apply) every time may be a bit wasteful, since we know it can only be None or that. Anyway at least matching x.owner.op: case Elemwise(scalar_op=ps.Add()) would already be a bit more readable and no less performant

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