flax
c5f2e5b0 - Make Flax thread-safe for multi-threaded tracing.

Commit
3 days ago
Make Flax thread-safe for multi-threaded tracing. Replace the global _use_named_call bool with a contextvars.ContextVar so each thread gets its own context without racing on a shared global. Move submodule adoption caching from a process-global WeakKeyDictionary onto flax.core.Scope (_submodule_adoption_cache) to guarantee thread-isolated tracing trees without global locks. PiperOrigin-RevId: 963430523
Author
Committer
Parents
Loading