jax

e20523c2 - Make api_test.py work when test cases are run using multiple threads.

Make api_test.py work when test cases are run using multiple threads. * keep track of all known config.State objects so we can find them by name. * change `@jtu.with_config` to default to setting thread-local configurations. * add a `@jtu.with_global_config` for those things that truly need to be set globally. * add a `@jtu.thread_local_config_context` that overrides thread-local configuration options, just as `jtu.global_config_context` overrides global configuration options. * change the pretty printer color option to be a State so it can be set locally. * tag a number of tests as thread-hostile, in particular tests that check counters for numbers of compilations, rely on garbage collection having particular semantics, or look at log output. PiperOrigin-RevId: 713411171