Cache used_axis_names calls for jaxprs
Cache used_axis_names calls for jaxprs

All control-flow primitives are `AxisPrimitive`s now, which means that we're doing
lots of those used-names traversals during dispatch and they can be expensive!
This adds caching to try to lower the cost of that change.

PiperOrigin-RevId: 395921887

google/jax,google/jax,google/jax,google/jax