provide an easier api for transpiling to jax for usage with optimistix

[kratsg]

./cc @pfackeldey

Follow up from the work done in pymc-devs/pytensor#1636

[kratsg]

see also this gist: https://gist.github.com/kratsg/2c83d4dd3367664554ba842e85c48879