Thanks yes that would work in that case, I wasn’t aware of the extra_ops.
Is there an alternative to this approach for more complicated calculation? I.e. should we always find a way to express the calculation in the form of a sequence of pytensor operations?