Fix the breakage caused by deleted enable_memories config

PiperOrigin-RevId: 707331603

Author: yashk2810

jax/_src/config.py

@@ -215,6 +215,7 @@ def trace_context():
           use_direct_linearize.value,
           softmax_custom_jvp.value,
           disable_jit.value,
+          enable_memories.value,
           debug_key_reuse.value,
           jax_xla_profile_version.value,
           # Technically this affects jaxpr->stablehlo lowering, not tracing.
@@ -971,10 +972,20 @@ def enum_flag(name, default, *args, **kwargs) -> Flag[str]:
     upgrade=True,
     help='If True, pmap and shard_map API will be merged.')
 
+def _update_jax_memories_global(val):
+  if hasattr(jax_jit.global_state(), 'enable_memories'):
+    jax_jit.global_state().enable_memories = val
+
+def _update_jax_memories_thread_local(val):
+  if hasattr(jax_jit.thread_local_state(), 'enable_memories'):
+    jax_jit.thread_local_state().enable_memories = val
+
 enable_memories = bool_state(
     'jax_enable_memories',
     default=True,
     upgrade=True,
+    update_global_hook=_update_jax_memories_global,
+    update_thread_local_hook=_update_jax_memories_thread_local,
     help=("If True, will allow fetching memory kinds available on executable "
           "and annotate Shardings with it."))