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: 963582520
