use unsafe jax.core function rather than jax.core internal data struc…

…tures so as to be slightly more robust to jax core changes PiperOrigin-RevId: 684632943
google-research · Oct 11, 2024 · c5a27ea · c5a27ea
1 parent 8f58121
commit c5a27ea
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/scenic/projects/lang4video/trainer/train_utils.py b/scenic/projects/lang4video/trainer/train_utils.py
@@ -221,8 +221,7 @@ def _fun(*args_, **kwargs_):
 
 
 def axis_name_exists(axis_name: Hashable) -> bool:
-  return any(frame.name == axis_name
-             for frame in jax.core.thread_local_state.trace_state.axis_env)
+  return axis_name in jax.core.unsafe_get_axis_names()  # type: ignore
 
 
 # The following are the same functions in `jax.example_libraries.optimizers` but