The as_op decorator seems unnecessary. You can use pytensor operations, and if you want to exploit numba for the the whole logp (and automatically dlogp because you didn’t use as_op), pass compile_kwargs=dict(backend="NUMBA") to pm.sample
For the intricacies of ODEs @aseyboldt may have some hints