Tips

Don’t define new Ops unless you have to

It is usually not useful to define Ops that can be easily implemented using other already existing Ops. For example, instead of writing a “sum_square_difference” Op, you should probably just write a simple function:

from pytensor import tensor as pt

def sum_square_difference(a, b):
    return pt.sum((a - b)**2)

Even without taking PyTensor’s rewrites into account, it is likely to work just as well as a custom implementation. It also supports all data types, tensors of all dimensions as well as broadcasting, whereas a custom implementation would probably only bother to support contiguous vectors/matrices of doubles…

Use PyTensor’s high order Ops when applicable

PyTensor provides some generic Op classes which allow you to generate a lot of Ops at a lesser effort. For instance, Elemwise can be used to make elemwise operations easily, whereas DimShuffle can be used to make transpose-like transformations. These higher order Ops are mostly tensor-related, as this is PyTensor’s specialty.