Open
Description
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