Implement helper @as_jax_op to wrap JAX functions in PyTensor

[ricardoV94]

Description

This blogpost walks through the logic for 3 different examples: https://github.com/pymc-labs/communications/issues/21 and shows the logic is always the same:

1. Wrap jitted forward pass in Op
2. Wrap jitted jvp (or vjp I can never remember) as a GradOp to provide gradient implementation
3. Dispatch unjitted versions of the two Ops for integration with `function(... , mode="JAX")

Things that cannot be obtained automatically (or maybe they can?) and should be opt-in as in @as_op:
4. Input and outputs types
5. infer_shape