jax

1e813e16 - [better_errors] Continue adding debug info to Jaxprs (step 4)

[better_errors] Continue adding debug info to Jaxprs (step 4) This follows after #26078, #26313, #26348, adding `debug_info` to more calls to `lu.wrap_init`. As part of this I have changed the primitive `custom_transpose` to take the `transpose` parameter as a `lu.WrappedFun`, which carries debug info. Previously, this was a `Callable`. These changes ensure that all the `lu.wrap_init` and `Jaxpr` are called with debug_info in the `api_test.py:CustomTransposeTest`.