jax

be1a554d - Fix a possible race in pjit.cc.

Fix a possible race in pjit.cc. We need to be careful not to destroy Python objects while using a Python 3.13- critical section to protect C++ state. The critical section might be released when calling back into Python code (much as the GIL may be released in GIL mode). In this code Key is kept alive by the function already, but the Value may be deleted before the hash table updates are done. PiperOrigin-RevId: 744008939