jax
9819417d - Add `config.trace_context_names` for better cache miss messages.

Commit
180 days ago
Add `config.trace_context_names` for better cache miss messages. This change introduces `jaxlib.config.trace_context_names`, which returns the names of the configuration values included in the JAX tracing context. This is used in `partial_eval.py` to provide more detailed explanations when a cache miss occurs due to differences in the tracing context, specifying which named config values have changed. The new functionality is guarded by a jaxlib version check. PiperOrigin-RevId: 917195862
Author
Parents
Loading