Add itypes/otypes specification for as_jax_op/wrap_jax #1618
Description
Description
Users should be able to set ityes and otypes for as_jax_op/wrap_jax directly if they want to avoid automatic inference, just as in as_op.
{{ message }}
Users should be able to set ityes and otypes for as_jax_op/wrap_jax directly if they want to avoid automatic inference, just as in as_op.