If you are familiar with JAX, the grad method (or even better the L_op) is very similar to the vjp: There are some examples of how they match here: How to use JAX ODEs and Neural Networks in PyMC - PyMC Labs
Otherwise, if you can share the exact Op you are trying to implement, we can provide more specific guidance.