cache tracing of (sub)calls when forming a jaxpr

[mattjj]

cc @hawkinsp

currently includes changes from #9188

The main motivations here are:

1. reduce trace times (by not retracing functions to a jaxpr when we can just look up the result in a cache), and
2. reduce compilation times by lowering outlined (as opposed to inlined) functions to our compiler targets.

This PR on its own tackles (1) but not (2). It sets us up for (2) because subjaxprs are now identical Python objects when we trace to a jaxpr, and so we can have have lowerings which are themselves memoized on (sub)jaxpr object id or alternatively we can have an early pass to produce a JaxprModule with explicit outlined functions. The latter seems more robust and clearer.

fixes #7155

TODO:

• traceback surgery for source info (that's the last failing test)

[mattjj]

@patrick-kidger check out this PR if you haven't already! Steps towards a non-inlined and faster-compiling world.

[mattjj]

I think this is subsumed by #10775 (for the subjaxpr stuff) and #11298 (for the axis env stuff).