You have to wrap your function in an Op, which the blackbox likelihood notebook demonstrates, there is no other way around.
Here are some examples with wrapping JAX functions that may feel more succinct: How to use JAX ODEs and Neural Networks in PyMC - PyMC Labs