jax

59e5ce22 - Avoid calls to warnings.catch_warnings in JAX core code.

Avoid calls to warnings.catch_warnings in JAX core code. warnings.catch_warnings is not thread-safe. However it is always used to avoid complex-to-real conversion warnings, which we can avoid in other ways.