code: def enable_jax_debugging_flags():
      """Enables some of the global JAX flags for debugging."""
  
      # Enable the NaN-checker behavior to cause JAX to hard-break on the first
      # occurrence of a NaN.
      jax.config.update("jax_debug_nans", True)
  
      # Enable the compilation logger to check whether or not we're accidentally
      # causing a lot of re-compilation (inspect logs for excessive jitting).
      jax.config.update("jax_log_compiles", True)
  
      # Detect numpy-style automatic rank promotion and force strict, explicit
      # casts. We can use `raise` instead of warn to raise an error.
      jax.config.update("jax_numpy_rank_promotion", "warn")
  
      # Print global JAX flags in logs.
      logging.info("Global JAX flags: %s", jax.config.values)
  
imports: from collections import abc
  from concurrent import futures
  import json
  import operator
  import threading
  from typing import Any, Callable, Optional, Sequence, Set, Tuple, Union
  from absl import logging
  from clu import parameter_overview
  import jax
  import jax.numpy
  from jax.tree_util import tree_map
  import ml_collections
label_desc: Use image node_id 1 for calling a global function with example usage: enable_jax_debugging_flags() without return types