fix: stacktraces for custom exception handling

[pfackeldey]

This PR fixes stack traces for awkward's custom error handling. It allows to use debuggers like "pdb" that try to go through the stack traces again, before they've been evicted from the stack (or wrongly ordered?) (I'm not sure what causes that exactly). This PR does not affect the custom decoration of errors.

[pfackeldey]

Could you have a look @henryiii on this PR? I remember you had some ideas about it in one of the recent awkward meetings

[jpivarski]

This looks good! (Minor comment below.)

The JAX test failures should be dealt with elsewhere and in one place, perhaps as a single PR that fixes or removes them and then gets merged into all of these pending PRs.

[pfackeldey]

Great! The jax test fix is in #3340 already :)

[henryiii]

I don't think it's related, but there's a lot of annotations like:

jax.errors.UnexpectedTracerError: Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type float64[3] wrapped in a JVPTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.
To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.UnexpectedTracerError

being produced.

Edit: Ah, I see Jim mentioned this. Edit: And Peter mentioned a fix. I'm behind. :)