jax
35fe85c3 - [trace_to_jaxpr] Create `trace_to_jaxpr_nocache`

Commit
19 days ago
[trace_to_jaxpr] Create `trace_to_jaxpr_nocache` It is sometimes clear that the cache on `trace_to_jaxpr` will miss, e.g., when we are tracing a locally created function somewhere in JAX internals. For those cases I think it is better to explicitly bypass the cache, because we would just miss, add an entry to the cache, perhaps print a useless cache miss explanation, possibly evicting something useful from there, and then we would immediately remove the newly added entry from the cache because the function goes out of scope. We are in the process of replacing uses of `pe.trace_to_jaxpr_dynamic` with uses of `pe.trace_to_jaxpr`. Note that `pe.trace_to_jaxpr_dynamic` was not even a cached function. So, having some of the replacements explicitly bypass the cache is not a regression. An example is [this use](https://github.com/jax-ml/jax/blob/452abf87bfb494efedc0be2d689a2b771a9731be/jax/_src/lax/control_flow/loops.py#L1314) but there are others, and there will be even more as we replace `trace_to_jaxpr_dynamic` This change only introduces the function, but it does not introduce uses of the new function. PiperOrigin-RevId: 935258398
Author
Parents
Loading