# Copyright 2018 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Callable, Iterator, Sequence
import contextlib
import enum
import functools
import itertools
import logging
import os
import sys
from typing import Any, Generic, NoReturn, Optional, Protocol, TypeVar, cast

from jax._src.lib import guard_lib
from jax._src.lib import jax_jit
from jax._src.lib import xla_client
from jax._src import logging_config

config_ext = xla_client._xla.config

logger = logging.getLogger(__name__)

_T = TypeVar('_T')


class EffortLevel(enum.Enum):
  """Effort level enum, mirroring the XLA effort options."""

  UNKNOWN = 0
  O0 = 9
  O1 = 19
  O2 = 29
  O3 = 39

  @classmethod
  def _missing_(cls, value: object) -> EffortLevel | None:
    return _effort_from_string.get(value)


_effort_from_string: dict[Any, EffortLevel] = {
    'UNKNOWN': EffortLevel.UNKNOWN,
    'O0': EffortLevel.O0,
    'O1': EffortLevel.O1,
    'O2': EffortLevel.O2,
    'O3': EffortLevel.O3,
}


def bool_env(varname: str, default: bool) -> bool:
  """Read an environment variable and interpret it as a boolean.

  True values are (case insensitive): 'y', 'yes', 't', 'true', 'on', and '1';
  false values are 'n', 'no', 'f', 'false', 'off', and '0'.

  Args:
    varname: the name of the variable
    default: the default boolean value
  Raises: ValueError if the environment variable is anything else.
  """
  val = os.getenv(varname, str(default))
  val = val.lower()
  if val in ('y', 'yes', 't', 'true', 'on', '1'):
    return True
  elif val in ('n', 'no', 'f', 'false', 'off', '0'):
    return False
  else:
    raise ValueError(f"invalid truth value {val!r} for environment {varname!r}")

def int_env(varname: str, default: int) -> int:
  """Read an environment variable and interpret it as an integer."""
  return int(os.getenv(varname, str(default)))


class ValueHolder(Protocol[_T]):
  """A holder for a configuration value.

  There are two kinds of value holders: ``Flag``, which is assigned exactly
  once and never modified after; and ``State``, which can be changed locally
  within a thread via a context manager.
  """

  value: _T

  def _set(self, value: _T) -> None: ...


class Config:
  _HAS_DYNAMIC_ATTRIBUTES = True

  def __init__(self):
    self._value_holders: dict[str, ValueHolder] = {}
    self.meta = {}
    self.use_absl = False
    self._contextmanager_flags = set()

  def update(self, name, val):
    if name not in self._value_holders:
      raise AttributeError(f"Unrecognized config option: {name}")
    self._value_holders[name]._set(val)

  def read(self, name):
    if name in self._contextmanager_flags:
      raise AttributeError(
          "For flags with a corresponding contextmanager, read their value "
          f"via e.g. `config.{name}` rather than `config.FLAGS.{name}`.")
    return self._read(name)

  def _read(self, name):
    try:
      return self._value_holders[name].value
    except KeyError:
      raise AttributeError(f"Unrecognized config option: {name}")

  @property
  def values(self):
    return {name: holder.value for name, holder in self._value_holders.items()}

  def add_option(self, name, holder, opt_type, meta_args, meta_kwargs):
    if name in self._value_holders:
      raise Exception(f"Config option {name} already defined")
    self._value_holders[name] = holder
    self.meta[name] = (opt_type, meta_args, meta_kwargs)

  def config_with_absl(self):
    """Registers absl flags for the JAX configs.

    E.g., for each JAX config defined using bool_state(), this method
    registers an absl boolean flag, with the same name.

    This is the recommended method to call if you use `app.run(main)` and you
    need JAX flags.

    Examples:

    ```python
    from absl import app
    import jax
    ...

    if __name__ == '__main__':
      jax.config.config_with_absl()
      app.run(main)
    ```

    """
    import absl.flags as absl_FLAGS  # noqa: F401  # pytype: disable=import-error
    from absl import app, flags as absl_flags  # pytype: disable=import-error

    self.use_absl = True
    self.absl_flags = absl_flags
    absl_defs = { bool: absl_flags.DEFINE_bool,
                  int:  absl_flags.DEFINE_integer,
                  float: absl_flags.DEFINE_float,
                  str:  absl_flags.DEFINE_string,
                  'enum': absl_flags.DEFINE_enum }

    for name, (flag_type, meta_args, meta_kwargs) in self.meta.items():
      holder = self._value_holders[name]
      absl_defs[flag_type](name, holder.value, *meta_args, **meta_kwargs)
    app.call_after_init(lambda: self.complete_absl_config(absl_flags))

  def complete_absl_config(self, absl_flags):
    # NOTE: avoid calling from outside this module. Instead, use
    # `config_with_absl()`, and (in rare cases) `parse_flags_with_absl()`.
    for name, holder in self._value_holders.items():
      try:
        flag = absl_flags.FLAGS[name]
      except KeyError:
        # This can happen if a new flag was added after config_with_absl() was
        # called, but before complete_absl_config was run. We could in principle
        # add code to DEFINE_... to register any newly added flags with ABSL
        # if config_with_absl() has already been called, but arguably the user
        # should have called config_with_absl() later.
        continue
      if flag.present:
        holder._set(flag.value)

  def parse_flags_with_absl(self):
    """Parses command-line args that start with --jax.

    This method should be used only by advanced users. Most users should use
    :meth:`config_with_absl` instead.

    This method has serious limitations: e.g., although it parses only the
    --jax* command-line args, it runs the validators of all registered absl
    flags, even non-JAX ones that have not been set yet; as such, for the
    non-JAX flags, the validators run on the default flag values, not on the
    values indicated by the command-line args.
    """
    global already_configured_with_absl
    if not already_configured_with_absl:
      # Extract just the --jax... flags (before the first --) from argv. In some
      # environments (e.g. ipython/colab) argv might be a mess of things
      # parseable by absl and other junk.
      jax_argv = itertools.takewhile(lambda a: a != '--', sys.argv)
      jax_argv = ['', *(a for a in jax_argv if a.startswith('--jax'))]

      import absl.flags  # pytype: disable=import-error
      self.config_with_absl()
      absl.flags.FLAGS(jax_argv, known_only=True)
      self.complete_absl_config(absl.flags)
      already_configured_with_absl = True


def trace_context():
  """Returns a tuple of configuration values that affect tracing.

  These values are included in the cache key for linear_util.cache.

  Values included in this set should also most likely be included in
  the C++ JIT state, which is handled separately.
  """
  return (axis_env_state.value, mesh_context_manager.value,
          xla_metadata_context_manager.value,
          abstract_mesh_context_manager.value,
          compute_on_context_manager.value, enable_x64.value,
          numpy_rank_promotion.value, default_matmul_precision.value,
          dynamic_shapes.value,
          eager_constant_folding.value,
          numpy_dtype_promotion.value,
          default_device.value, random_seed_offset.value,
          threefry_partitionable.value,
          threefry_gpu_kernel_lowering.value,
          use_direct_linearize.value,
          varying_axes_in_types.value,
          softmax_custom_jvp.value,
          disable_jit.value,
          debug_key_reuse.value,
          jax_xla_profile_version.value,
          # Technically this affects jaxpr->stablehlo lowering, not tracing.
          hlo_source_file_canonicalization_regex.value,
          pgle_profiling_runs.value,
          enable_pgle.value,
          use_shardy_partitioner.value,
          use_high_dynamic_range_gumbel.value)

config = Config()

_read = config._read
update = config.update
parse_flags_with_absl = config.parse_flags_with_absl


class NoDefault: pass
no_default = NoDefault()

config_states = {}

class State(config_ext.Config[_T]):

  __slots__ = (
      '_name', '_update_thread_local_hook', '_update_global_hook',
      '_validator', '_default_context_manager_value', '__doc__', '__name__',
  )

  def __init__(
      self,
      name: str,
      default: _T,
      help,
      update_global_hook: Callable[[_T], None] | None = None,
      update_thread_local_hook: Callable[[_T | None], None] | None = None,
      validator: Callable[[Any], None] | None = None,
      extra_description: str = '',
      default_context_manager_value: Any = no_default,
      include_in_jit_key: bool = False,
  ):
    super().__init__(default, include_in_jit_key)
    self._name = name
    self.__name__ = name[4:] if name.startswith('jax_') else name
    self.__doc__ = (f"Context manager for `{name}` config option"
                    f"{extra_description}.\n\n{help}")
    self._update_global_hook = update_global_hook
    self._update_thread_local_hook = update_thread_local_hook
    self._validator = validator
    self._default_context_manager_value = default_context_manager_value
    if self._validator:
      self._validator(default)
    if self._update_global_hook:
      self._update_global_hook(default)
    config_states[name] = self

  def __bool__(self) -> NoReturn:
    raise TypeError(
        "bool() not supported for instances of type '{0}' "
        "(did you mean to use '{0}.value' instead?)".format(
            type(self).__name__))

  def _set(self, value: _T) -> None:
    if self._validator:
      self._validator(value)
    self.set_global(value)
    if self._update_global_hook:
      self._update_global_hook(value)

  def __call__(self, new_val: Any = no_default):
    return StateContextManager(self, new_val)

  def _add_hooks(self, update_global_hook, update_thread_local_hook):
    """Private method that adds hooks to an existing context-manager.

    Used to avoid cyclic import dependencies."""
    self._update_thread_local_hook = update_thread_local_hook
    self._update_global_hook = update_global_hook
    update_global_hook(self.get_global())


class StateContextManager(contextlib.ContextDecorator):
  __slots__ = ['state', 'new_val', 'prev']

  def __init__(self, state, new_val):
    self.state = state
    self.new_val = new_val

    if new_val is no_default:
      if state._default_context_manager_value is not no_default:
        new_val = state._default_context_manager_value  # default_context_manager_value provided to constructor
      else:
        # no default_value provided to constructor and no value provided as an
        # argument, so we raise an error
        raise TypeError(f"Context manager for {state.__name__} config option "
                        "requires an argument representing the new value for "
                        "the config option.")
    if state._validator:
      state._validator(new_val)


  def __enter__(self):
    self.prev = self.state.swap_local(self.new_val)
    if self.state._update_thread_local_hook:
      self.state._update_thread_local_hook(self.new_val)

  def __exit__(self, exc_type, exc_value, traceback):
    self.state.set_local(self.prev)
    if self.state._update_thread_local_hook:
      if self.prev is config_ext.unset:
        self.state._update_thread_local_hook(None)
      else:
        self.state._update_thread_local_hook(cast(Optional[Any], self.prev))


UPGRADE_BOOL_HELP = (
    " This will be enabled by default in future versions of JAX, at which "
    "point all uses of the flag will be considered deprecated (following "
    "the `API compatibility policy "
    "<https://jax.readthedocs.io/en/latest/api_compatibility.html>`_).")

UPGRADE_BOOL_EXTRA_DESC = " (transient)"


def bool_state(
    name: str,
    default: bool,
    help: str,
    *,
    update_global_hook: Callable[[bool], None] | None = None,
    update_thread_local_hook: Callable[[bool | None], None] | None = None,
    upgrade: bool = False,
    extra_description: str = '',
    include_in_jit_key: bool = False,
) -> State[bool]:
  """Set up thread-local state and return a contextmanager for managing it.

  This function is a convenience wrapper. It defines a flag, environment
  variable, and corresponding thread-local state, which can be managed via the
  contextmanager it returns.

  The thread-local state value can be read via the ``config.<option_name>``
  attribute, where ``config`` is the singleton ``Config`` instance.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    default: boolean, a default value for the option.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.
    update_global_hook: a optional callback that is called with the updated
      value of the global state when it is altered or set initially.
    update_thread_local_hook: a optional callback that is called with the
      updated value of the thread-local state when it is altered or set
      initially.
    upgrade: optional indicator that this flag controls a canonical feature
      upgrade, so that it is `True` for the incoming functionality, `False`
      for the outgoing functionality to be deprecated.
    extra_description: string, optional: extra information to add to the
      summary description.

  Returns:
    A contextmanager to control the thread-local state value.

  Examples:

    ENABLE_FOO = config.bool_state(
        name='jax_enable_foo',
        default=False,
        help='Enable foo.')

    # Now the JAX_ENABLE_FOO shell environment variable and --jax_enable_foo
    # command-line flag can be used to control the process-level value of
    # the configuration option, in addition to using e.g.
    # ``config.update("jax_enable_foo", True)`` directly. We can also use a
    # context manager:

    with enable_foo(True):
      ...

  The value of the thread-local state or flag can be accessed via
  ``config.jax_enable_foo``. Reading it via ``config.FLAGS.jax_enable_foo`` is
  an error.
  """
  if not isinstance(default, bool):
    raise TypeError(f"Default value must be of type bool, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  default = bool_env(name.upper(), default)
  name = name.lower()
  if upgrade:
    help += ' ' + UPGRADE_BOOL_HELP
    extra_description += UPGRADE_BOOL_EXTRA_DESC
  config._contextmanager_flags.add(name)

  s = State[bool](
      name, default, help, update_global_hook=update_global_hook,
      update_thread_local_hook=update_thread_local_hook,
      extra_description=extra_description, default_context_manager_value=True,
      include_in_jit_key=include_in_jit_key)
  config.add_option(name, s, bool, meta_args=[], meta_kwargs={"help": help})
  setattr(Config, name, property(lambda _: s.value))
  return s


def enum_state(
    name: str,
    enum_values: Sequence[str],
    default: str,
    help: str,
    *,
    update_global_hook: Callable[[str], None] | None = None,
    update_thread_local_hook: Callable[[str | None], None] | None = None,
    include_in_jit_key: bool = False,
) -> State[str]:
  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    enum_values: list of strings representing the possible values for the
      option.
    default: string, default value.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if not isinstance(default, str):
    raise TypeError(f"Default value must be of type str, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  name = name.lower()
  default = os.getenv(name.upper(), default)
  if default not in enum_values:
    raise ValueError(f"Invalid value \"{default}\" for JAX flag {name}")
  config._contextmanager_flags.add(name)

  def validator(new_val):
    if type(new_val) is not str or new_val not in enum_values:
      raise ValueError(f"new enum value must be in {enum_values}, "
                       f"got {new_val} of type {type(new_val)}.")

  s = State[str](
      name,
      default,
      help,
      update_global_hook=update_global_hook,
      update_thread_local_hook=update_thread_local_hook,
      validator=validator,
      include_in_jit_key=include_in_jit_key,
  )
  config.add_option(
      name, s, 'enum',
      meta_args=[],
      meta_kwargs={"enum_values": enum_values, "help": help}
  )
  setattr(Config, name, property(lambda _: s.value))
  return s


def optional_enum_state(
    name: str,
    enum_values: Sequence[str],
    default: str | None,
    help: str,
    *,
    update_global_hook: Callable[[str | None], None] | None = None,
    update_thread_local_hook: Callable[[str | None], None] | None = None,
    include_in_jit_key: bool = False,
) -> State[str | None]:
  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    enum_values: list of strings representing the possible values for the
      option.
    default: optional string, default value.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if default is not None and not isinstance(default, str):
    raise TypeError(f"Default value must be of type str or None, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  name = name.lower()
  default = os.getenv(name.upper(), default)
  if default is not None and default not in enum_values:
    raise ValueError(f"Invalid value \"{default}\" for JAX flag {name}")
  config._contextmanager_flags.add(name)

  def validate(new_val):
    if (new_val is not None and
      (type(new_val) is not str or new_val not in enum_values)):
      raise ValueError(f"new enum value must be None or in {enum_values}, "
                       f"got {new_val} of type {type(new_val)}.")

  s = State['str | None'](
      name, default, help, update_global_hook, update_thread_local_hook,
      validate, include_in_jit_key=include_in_jit_key,
  )
  config.add_option(
      name, s, 'enum',
      meta_args=[],
      meta_kwargs={"enum_values": enum_values, "help": help}
  )
  setattr(Config, name, property(lambda _: s.value))
  return s


def int_state(
    name: str,
    default: int,
    help: str,
    *,
    update_global_hook: Callable[[int], None] | None = None,
    update_thread_local_hook: Callable[[int | None], None] | None = None,
    include_in_jit_key: bool = False,
) -> State[int]:
  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    default: optional int, default value.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if not isinstance(default, int):
    raise TypeError(f"Default value must be of type int, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  name = name.lower()
  default_env = os.getenv(name.upper())
  if default_env is not None:
    try:
      default = int(default_env)
    except ValueError:
      raise ValueError(f"Invalid value \"{default_env}\" for JAX flag {name}")
  config._contextmanager_flags.add(name)

  def validate(new_val):
    if new_val is not None and not isinstance(new_val, int):
      raise ValueError(f'new int config value must be None or of type int, '
                       f'got {new_val} of type {type(new_val)}')

  s = State[int](name, default, help, update_global_hook,
                 update_thread_local_hook, validate,
                 include_in_jit_key=include_in_jit_key)
  config.add_option(name, s, int, meta_args=[], meta_kwargs={"help": help})
  setattr(Config, name, property(lambda _: s.value))
  return s


def float_state(
    name: str,
    default: float,
    help: str,
    *,
    update_global_hook: Callable[[float], None] | None = None,
    update_thread_local_hook: Callable[[float | None], None] | None = None,
) -> State[float]:
  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    default: default value.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if not isinstance(default, float):
    raise TypeError(f"Default value must be of type float, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  name = name.lower()
  default_env = os.getenv(name.upper())
  if default_env is not None:
    try:
      default = float(default_env)
    except ValueError:
      raise ValueError(f"Invalid value \"{default_env}\" for JAX flag {name}")
  config._contextmanager_flags.add(name)

  def validate(new_val):
    if new_val is not None and not isinstance(new_val, (float, int)):
      raise ValueError(
        f'new float config value must be None or of type float, '
        f'got {new_val} of type {type(new_val)}')

  s = State[float](name, default, help, update_global_hook,
                   update_thread_local_hook, validate)
  config.add_option(name, s, float, meta_args=[], meta_kwargs={"help": help})
  setattr(Config, name, property(lambda _: s.value))
  return s


def string_state(
    name: str,
    default: str,
    help: str,
    *,
    update_global_hook: Callable[[str], None] | None = None,
    update_thread_local_hook: Callable[[str | None], None] | None = None,
) -> State[str]:
  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    default: string, a default value for the option.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.
    update_global_hook: an optional callback that is called with the updated
      value of the global state when it is altered or set initially.
    update_thread_local_hook: an optional callback that is called with the
      updated value of the thread-local state when it is altered or set
      initially.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if not isinstance(default, str):
    raise TypeError(f"Default value must be of type str, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")

  def validator(new_val):
    if not isinstance(new_val, str):
      raise TypeError('new string config value must be of type str,'
                       f' got {new_val} of type {type(new_val)}.')

  return string_or_object_state(
      name, default, help,
      update_global_hook=update_global_hook,
      update_thread_local_hook=update_thread_local_hook,
      validator=validator)


def optional_string_state(
    name: str,
    default: str | None,
    help: str,
    *,
    update_global_hook: Callable[[str], None] | None = None,
    update_thread_local_hook: Callable[[str | None], None] | None = None,
) -> State[str | None]:
  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    default: optional string, a default value for the option.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.
    update_global_hook: an optional callback that is called with the updated
      value of the global state when it is altered or set initially.
    update_thread_local_hook: an optional callback that is called with the
      updated value of the thread-local state when it is altered or set
      initially.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if default is not None and not isinstance(default, str):
    raise TypeError(f"Default value must be of type str or None, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")

  def validator(new_val):
    if new_val is not None and not isinstance(new_val, str):
      raise ValueError('new string config value must be None or of type str,'
                       f' got {new_val} of type {type(new_val)}.')

  return string_or_object_state(
      name, default, help,
      update_global_hook=update_global_hook,
      update_thread_local_hook=update_thread_local_hook,
      validator=validator)

def string_or_object_state(
    name: str,
    default: Any,
    help: str,
    *,
    update_global_hook: Callable[[Any], None] | None = None,
    update_thread_local_hook: Callable[[Any], None] | None = None,
    validator: Callable[[Any], None] | None = None,
) -> State[Any]:
  """Set up thread-local state and return a contextmanager for managing it.

  Similar to ``string_state``, except the context manager will accept
  any object, not just a string. Any value passed via command line flag or
  environment variable will be treated as a string.

  Args:
    name: string, converted to lowercase to define the name of the config
      option (and absl flag). It is converted to uppercase to define the
      corresponding shell environment variable.
    default: string, a default value for the option.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.
    update_global_hook: an optional callback that is called with the updated
      value of the global state when it is altered or set initially.
    update_thread_local_hook: an optional callback that is called with the
      updated value of the thread-local state when it is altered or set
      initially.
    validator: an optional callback that is called with the new
      value on any update, and should raise an error if the new value is
      invalid.

  Returns:
    A contextmanager to control the thread-local state value.
  """
  name = name.lower()
  default = os.getenv(name.upper(), default)
  config._contextmanager_flags.add(name)

  s = State[Any](
      name, default, help, update_global_hook, update_thread_local_hook,
      validator)
  setattr(Config, name, property(lambda _: s.value))
  config.add_option(name, s, str, meta_args=[], meta_kwargs={"help": help})
  return s


class Flag(Generic[_T]):

  __slots__ = ("_name", "value", "_update_hook")

  _name: str
  value: _T
  _update_hook: Callable[[Any], None] | None

  def __init__(self, name: str, default: _T,
               update_hook: Callable[[Any], None] | None = None):
    self._name = name
    self._update_hook = update_hook
    self._set(default)

  def __bool__(self) -> NoReturn:
    raise TypeError(
        "bool() not supported for instances of type '{0}' "
        "(did you mean to use '{0}.value' instead?)".format(
            type(self).__name__))

  def _set(self, value: _T) -> None:
    self.value = value
    if self._update_hook is not None:
      self._update_hook(value)


def bool_flag(name, default, *args, **kwargs) -> Flag[bool]:
  update_hook = kwargs.pop("update_hook", None)
  holder = Flag(name, default, update_hook)
  config.add_option(name, holder, bool, args, kwargs)
  return holder


def int_flag(name, default, *args, **kwargs) -> Flag[int]:
  update_hook = kwargs.pop("update_hook", None)
  holder = Flag(name, default, update_hook)
  config.add_option(name, holder, int, args, kwargs)
  return holder


def float_flag(name, default, *args, **kwargs) -> Flag[float]:
  update_hook = kwargs.pop("update_hook", None)
  holder = Flag(name, default, update_hook)
  config.add_option(name, holder, float, args, kwargs)
  return holder


def string_flag(name, default, *args, **kwargs) -> Flag[str]:
  update_hook = kwargs.pop("update_hook", None)
  holder = Flag(name, default, update_hook)
  config.add_option(name, holder, str, args, kwargs)
  return holder


def enum_flag(name, default, *args, **kwargs) -> Flag[str]:
  update_hook = kwargs.pop("update_hook", None)
  holder = Flag(name, default, update_hook)
  config.add_option(name, holder, 'enum', args, kwargs)
  return holder


already_configured_with_absl = False


trace_state = config_ext.Config(None, include_in_jit_key=True)
axis_env_state = config_ext.Config((), include_in_jit_key=True)
mesh_context_manager = config_ext.Config((), include_in_jit_key=True)
abstract_mesh_context_manager = config_ext.Config(None, include_in_jit_key=True)
device_context = config_ext.Config(None, include_in_jit_key=True)
compute_on_context_manager = config_ext.Config(None, include_in_jit_key=True)
xla_metadata_context_manager = config_ext.Config(None, include_in_jit_key=True)


# TODO(b/214340779): remove flag when XLA:CPU is improved.
jax2tf_associative_scan_reductions = bool_state(
    name='jax2tf_associative_scan_reductions',
    default=False,
    help=(
        'JAX has two separate lowering rules for the cumulative reduction '
        'primitives (cumsum, cumprod, cummax, cummin). On CPUs and GPUs it uses '
        'a lax.associative_scan, while for TPUs it uses the HLO ReduceWindow. '
        'The latter has a slow implementation on CPUs and GPUs. '
        'By default, jax2tf uses the TPU lowering. Set this flag to True to '
        'use the associative scan lowering usage, and only if it makes a difference '
        'for your application. '
        'See the jax2tf README.md for more details.'
    )
)

jax2tf_default_native_serialization = bool_state(
    name='jax2tf_default_native_serialization',
    default=bool_env('JAX2TF_DEFAULT_NATIVE_SERIALIZATION', True),
    help=(
        'Sets the default value of the native_serialization parameter to '
        'jax2tf.convert. Prefer using the parameter instead of the flag, '
        'the flag may be removed in the future. '
        'Starting with JAX 0.4.31 non-native serialization is deprecated.'
    )
)

jax_serialization_version = int_state(
    name='jax_serialization_version',
    default=int_env('JAX_SERIALIZATION_VERSION', 0),  # We use 0 to detect default.
    help=(
        'DEPRECATED: use jax_export_calling_convention_version.'
    )
)

jax_export_calling_convention_version = int_state(
    name='jax_export_calling_convention_version',
    # Note: bump the default calling convention version at least one month after
    # we update XlaCallModule to support the new version, so that serialized
    # modules are forward compatible with deployed versions of XlaCallModule.
    # Version 9 of XlaCallModule is supported since October 27th, 2023.
    default=int_env('JAX_EXPORT_CALLING_CONVENTION_VERSION', 9),
    help=(
        'The calling convention version number to use for exporting. This must be '
        'within the range of versions supported by the tf.XlaCallModule '
        'used in your deployment environment. '
        'See https://jax.readthedocs.io/en/latest/export/shape_poly.html#calling-convention-versions.'
    )
)

export_ignore_forward_compatibility = bool_state(
    name='jax_export_ignore_forward_compatibility',
    default=bool_env('JAX_EXPORT_IGNORE_FORWARD_COMPATIBILIY', False),
    help=(
        'Whether to ignore the forward compatibility lowering rules. '
        'See https://jax.readthedocs.io/en/latest/export/export.html#compatibility-guarantees-for-custom-calls.'
    )
)

jax_platforms = optional_string_state(
    name='jax_platforms',
    default=None,
    help=(
        'Comma-separated list of platform names specifying which platforms jax '
        'should initialize. If any of the platforms in this list are not successfully '
        'initialized, an exception will be raised and the program will be aborted. '
        'The first platform in the list will be the default platform. '
        'For example, config.jax_platforms=cpu,tpu means that CPU and TPU backends '
        'will be initialized, and the CPU backend will be used unless otherwise '
        'specified. If TPU initialization fails, it will raise an exception. '
        'By default, jax will try to initialize all available '
        'platforms and will default to GPU or TPU if available, and fallback to CPU '
        'otherwise.'
        ))

jax_pjrt_client_create_options = optional_string_state(
    name='jax_pjrt_client_create_options',
    default=None,
    help=('A set of key-value pairs in the format of "k1:v1;k2:v2" strings '
          'provided to a device platform pjrt client as extra arguments.'))

enable_checks = bool_state(
    name='jax_enable_checks',
    default=False,
    help='Turn on invariant checking for JAX internals. Makes things slower.')

debug_key_reuse = bool_state(
    name='jax_debug_key_reuse',
    default=False,
    help=('Turn on experimental key reuse checking. With this configuration enabled,'
          ' typed PRNG keys (i.e. keys created with jax.random.key()) will have their'
          ' usage tracked, and incorrect reuse of a previously-used key will lead to'
          ' an error. Currently enabling this leads to a small Python overhead on'
          ' every call to a JIT-compiled function with keys as inputs or outputs.'))

check_tracer_leaks = bool_state(
    name='jax_check_tracer_leaks',
    default=False,
    help=('Turn on checking for leaked tracers as soon as a trace completes. '
          'Enabling leak checking may have performance impacts: some caching '
          'is disabled, and other overheads may be added. Additionally, be aware '
          'that some Python debuggers can cause false positives, so it is recommended '
          'to disable any debuggers while leak checking is enabled.'))
checking_leaks = functools.partial(check_tracer_leaks, True)

debug_nans = bool_state(
    name='jax_debug_nans',
    default=False,
    help=('Add nan checks to every operation. When a nan is detected on the '
          'output of a jit-compiled computation, call into the un-compiled '
          'version in an attempt to more precisely identify the operation '
          'which produced the nan.'))

debug_infs = bool_state(
    name='jax_debug_infs',
    default=False,
    help=('Add inf checks to every operation. When an inf is detected on the '
          'output of a jit-compiled computation, call into the un-compiled '
          'version in an attempt to more precisely identify the operation '
          'which produced the inf.'))

log_compiles = bool_state(
    name='jax_log_compiles',
    default=False,
    help=('Log a message each time `jit` or `pmap` compiles an XLA '
          'computation. Logging is performed with `logging`. When this '
          'option is set, the log level is WARNING; otherwise the level is '
          'DEBUG.'))

explain_cache_misses = bool_state(
    name='jax_explain_cache_misses',
    default=False,
    help=('Each time there is a miss on one of the main caches (e.g. the '
          'tracing cache), log an explanation.. Logging is performed with '
          '`logging`. When this option is set, the log level is WARNING; '
          'otherwise the level is DEBUG.'))

log_checkpoint_residuals = bool_state(
    name='jax_log_checkpoint_residuals',
    default=False,
    help=('Log a message every time jax.checkpoint (aka jax.remat) is '
          'partially evaluated (e.g. for autodiff), printing what residuals '
          'are saved.'))

pmap_shmap_merge = bool_state(
    name='jax_pmap_shmap_merge',
    default=False,
    upgrade=True,
    help='If True, pmap and shard_map API will be merged.')


spmd_mode = enum_state(
    name='jax_spmd_mode',
    enum_values=['allow_all', 'allow_jit'],
    default='allow_jit',
    help=("Decides whether Math on `jax.Array`'s that are not fully addressable "
          "(i.e. spans across multiple processes) is allowed. The options are: "
          "* allow_jit: Default, `pjit` and `jax.jit` computations are allowed "
          "    to execute on non-fully addressable `jax.Array`s\n"
          "* allow_all: `jnp`, normal math (like `a + b`, etc), `pjit`, "
          "    `jax.jit` and all other operations are allowed to "
          "    execute on non-fully addressable `jax.Array`s."))


distributed_debug = bool_state(
    name='jax_distributed_debug',
    default=False,
    help=('Enable logging useful for debugging multi-process distributed '
          'computations. Logging is performed with `logging` at WARNING '
          'level.'))

random_seed_offset = int_state(
    name='jax_random_seed_offset',
    default=0,
    help=('Offset to all random seeds (e.g. argument to jax.random.key()).'),
    include_in_jit_key=True,
)

legacy_prng_key = enum_state(
    name='jax_legacy_prng_key',
    enum_values=['allow', 'warn', 'error'],
    default='allow',
    help=('Specify the behavior when raw PRNG keys are passed to '
          'jax.random APIs.')
)

enable_custom_prng = bool_state(
    name='jax_enable_custom_prng',
    default=False,
    upgrade=True,
    help=('Enables an internal upgrade that allows one to define custom '
          'pseudo-random number generator implementations.'))

default_prng_impl = enum_state(
    name='jax_default_prng_impl',
    enum_values=['threefry2x32', 'rbg', 'unsafe_rbg'],
    default='threefry2x32',
    help=('Select the default PRNG implementation, used when one is not '
          'explicitly provided at seeding time.'))

threefry_partitionable = bool_state(
    name='jax_threefry_partitionable',
    default=True,
    upgrade=True,
    help=('Enables internal threefry PRNG implementation changes that '
          'render it automatically partitionable in some cases. Without this '
          'flag, using the standard jax.random pseudo-random number generation '
          'may result in extraneous communication and/or redundant distributed '
          'computation. With this flag, the communication overheads disappear '
          'in some cases.'),
    include_in_jit_key=True)

threefry_gpu_kernel_lowering = bool_state(
    name='jax_threefry_gpu_kernel_lowering',
    default=False,
    help=('On GPU, lower threefry PRNG operations to a kernel implementation. '
          'This makes compile times faster at a potential runtime memory '
          'cost.'),
    include_in_jit_key=True)

use_direct_linearize = bool_state(
    name='jax_use_direct_linearize',
    default=False,
    help=('Use direct linearization instead JVP followed by partial eval'),
    include_in_jit_key=True)

varying_axes_in_types = bool_state(
    name='jax_varying_axes_in_types',
    default=False,
    help=('Adds varying manual axes to ShapedArray to track which mesh axes the'
          ' array is varying over. This will help to remove the efficient'
          ' transpose rewrite machinery in shard_map'),
    include_in_jit_key=True)

data_dependent_tracing_fallback = bool_state(
    name='jax_data_dependent_tracing_fallback',
    default=False,
    help=('When True, falls back to trace dispatch based on data dependence '
          'instead of throwing an escaped tracer error.'))

softmax_custom_jvp = bool_state(
    name='jax_softmax_custom_jvp',
    default=False,
    upgrade=True,
    help=('Use a new custom_jvp rule for jax.nn.softmax. The new rule should '
          'improve memory usage and stability. Set True to use new '
          'behavior. See https://github.com/jax-ml/jax/pull/15677'),
    include_in_jit_key=True)


enable_custom_vjp_by_custom_transpose = bool_state(
    name='jax_enable_custom_vjp_by_custom_transpose',
    default=False,
    upgrade=True,
    help=('Enables an internal upgrade that implements `jax.custom_vjp` by '
          'reduction to `jax.custom_jvp` and `jax.custom_transpose`.'))

raise_persistent_cache_errors = bool_state(
    name='jax_raise_persistent_cache_errors',
    default=False,
    help=('If true, exceptions raised when reading or writing to the '
          'persistent compilation cache will be allowed through, halting '
          'program execution if not manually caught. If false, exceptions are '
          'caught and raised as warnings, allowing program execution to '
          'continue. Defaults to false so cache bugs or intermittent issues '
          'are non-fatal.'))

persistent_cache_min_compile_time_secs = float_state(
    name='jax_persistent_cache_min_compile_time_secs',
    default=1.,
    help=('The minimum compile time of a computation to be written to the '
          'persistent compilation cache. This threshold can be raised to '
          'decrease the number of entries written to the cache.'))

persistent_cache_min_entry_size_bytes = int_state(
    name='jax_persistent_cache_min_entry_size_bytes',
    default=0,
    help=('The minimum size (in bytes) of an entry that will be cached in the '
          'persistent compilation cache: '
          '* -1: disable the size restriction and prevent overrides. '
          '* Leave at default (0) to allow for overrides. The override will '
          '  typically ensure that the minimum size is optimal for the '
          '  filesystem being used for the cache. '
          '* > 0: the actual minimum size desired; no overrides.'))

# TODO: Change default to all
persistent_cache_enable_xla_caches = optional_string_state(
    name='jax_persistent_cache_enable_xla_caches',
    default='xla_gpu_per_fusion_autotune_cache_dir',
    help=('When the persistent cache is enabled, additional XLA caching will '
          'also be enabled automatically. This option can be used to configure'
          'which XLA caching methods will be enabled.'),
)

compilation_cache_include_metadata_in_key = bool_state(
    name='jax_compilation_cache_include_metadata_in_key',
    default=False,
    help=(
        'Include metadata, such as file names and line numbers, in the'
        ' compilation cache key. If false, the cache will still get hits even'
        ' if functions or files are moved, etc. However, it means that'
        ' executables loaded from the cache may have stale metadata, which'
        ' may show up in, e.g., profiles.'
    ),
)

hlo_source_file_canonicalization_regex = optional_string_state(
    name='jax_hlo_source_file_canonicalization_regex',
    default=None,
    help=('Used to canonicalize the source_path metadata of HLO instructions '
          'by removing the given regex. If set, re.sub() is called on each '
          'source_file with the given regex, and all matches are removed. '
          'This can be used to avoid spurious cache misses when using the '
          'persistent compilation cache, which includes HLO metadata in the '
          'cache key.'))

include_full_tracebacks_in_locations = bool_state(
    name='jax_include_full_tracebacks_in_locations',
    default=True,
    help=(
        'Include Python tracebacks in MLIR locations in IR emitted by JAX.'
    ),
)

traceback_in_locations_limit = int_state(
    name='jax_traceback_in_locations_limit',
    default=10,
    help=(
        'Limit the number of frames at the Python traceback frames included in '
        'MLIR locations. If set to the negative value, traceback will not be '
        'limited.'
    ),
)

share_binary_between_hosts = bool_state(
    name='jax_share_binary_between_hosts',
    default=False,
    help=(
        'If set to True, the compiled module will be shared between hosts '
        'directly.'
    ),
)

share_binary_between_hosts_timeout_ms = int_state(
    name='jax_share_binary_between_hosts_timeout_ms',
    default=20 * 60 * 1000,
    help='Timeout for the compiled module share.',
)

enable_pgle = bool_state(
    name='jax_enable_pgle',
    default=False,
    help=(
      'If set to True and the property jax_pgle_profiling_runs is set to '
      'greater than 0, the modules will be recompiled after running specified '
      'number times with collected data provided to the profile guided latency '
      'estimator.'
    ),
    include_in_jit_key=True,
)

pgle_profiling_runs = int_state(
    name='jax_pgle_profiling_runs',
    default=3,
    help=(
        'Amount of times module should be profiled before recompilation when '
        'PGLE is used.'
    ),
    include_in_jit_key=True,
)

pgle_aggregation_percentile = int_state(
    name='jax_pgle_aggregation_percentile',
    default=90,
    help='Percentile used to aggregate performance data between devices when '
         'PGLE is used.',
)

enable_compilation_cache = bool_state(
    name='jax_enable_compilation_cache',
    default=True,
    help=('If set to False, the compilation cache will be disabled regardless '
          'of whether set_cache_dir() was called. If set to True, the '
          'path could be set to a default value or via a call to '
          'set_cache_dir().'),
)

compilation_cache_dir = optional_string_state(
    name='jax_compilation_cache_dir',
    default=None,
    help=('Path for the cache. '
          'Precedence: '
          '1. A call to compilation_cache.set_cache_dir(). '
          '2. The value of this flag set in the command line or by default.'),
)

compilation_cache_expect_pgle = bool_state(
    name='jax_compilation_cache_expect_pgle',
    default=False,
    help=('If set to True, compilation cache entries that were compiled with '
          'profile data (i.e. PGLE was enabled and the requisite number of '
          'executions were profiled) will be preferentially loaded, even if '
          'PGLE is not currently enabled. A warning will be printed when no '
          'preferred cache entry is found.')
)

compilation_cache_max_size = int_state(
    name='jax_compilation_cache_max_size',
    default=-1,
    help=('The maximum size (in bytes) allowed for the persistent compilation '
          'cache. When set, the least recently accessed cache entry(s) '
          'will be deleted once the total cache directory size '
          'exceeds the specified limit. '
          'Caching will be disabled if this value is set to 0. A '
          'special value of -1 indicates no limit, allowing the cache '
          'size to grow indefinitely.'),
)

remove_custom_partitioning_ptr_from_cache_key = bool_state(
    name='jax_remove_custom_partitioning_ptr_from_cache_key',
    default=False,
    help=('If set to True, remove the custom partitioning pointer '
          'present in the precompiled stableHLO before hashing  '
          'during cache key computation. This is a potentially '
          'unsafe flag to set and only users who are sure of '
          'what they are trying to achieve should set it.'),
)

default_dtype_bits = enum_state(
    name='jax_default_dtype_bits',
    enum_values=['32', '64'],
    default='64',
    help=('Specify bit width of default dtypes, either 32-bit or 64-bit. '
          'This is a temporary flag that will be used during the process '
          'of deprecating the ``jax_enable_x64`` flag.'))

numpy_dtype_promotion = enum_state(
    name='jax_numpy_dtype_promotion',
    enum_values=['standard', 'strict'],
    default='standard',
    help=('Specify the rules used for implicit type promotion in operations '
          'between arrays. Options are "standard" or "strict"; in strict-mode, '
          'binary operations between arrays of differing strongly-specified '
          'dtypes will result in an error.'),
    include_in_jit_key=True)

disallow_mesh_context_manager = bool_state(
    name='jax_disallow_mesh_context_manager',
    default=False,
    help=(
        'If set to True, trying to use a mesh as a context manager will'
        ' result in a RuntimeError.'
    ),
)

def _update_x64_global(val):
  jax_jit.global_state().enable_x64 = val

def _update_x64_thread_local(val):
  jax_jit.thread_local_state().enable_x64 = val

enable_x64 = bool_state(
    name='jax_enable_x64',
    default=False,
    help='Enable 64-bit types to be used',
    update_global_hook=_update_x64_global,
    update_thread_local_hook=_update_x64_thread_local)

# TODO(phawkins): remove after fixing users of FLAGS.x64_enabled.
config._contextmanager_flags.remove('jax_enable_x64')

setattr(Config, "x64_enabled", property(lambda _: enable_x64.value))

def _update_default_device_global(val):
  jax_jit.global_state().default_device = val


def _update_default_device_thread_local(val):
  jax_jit.thread_local_state().default_device = val


def _validate_default_device(val):
  if (val is not None and
      not isinstance(val, xla_client.Device) and
      val not in ['cpu', 'gpu', 'tpu']):
    # TODO(skyewm): this is a workaround for non-PJRT Device types. Remove when
    # all JAX backends use a single C++ device interface.
    if 'Device' in str(type(val)):
      logger.info(
          'Allowing non-`xla_client.Device` default device: %s, type: %s',
          repr(val), type(val))
      return
    raise ValueError('jax.default_device must be passed either a Device object (e.g. '
                     f"`jax.devices('cpu')[0]`) or a platform name string like 'cpu' or 'gpu'"
                     f", got: {val!r}")


default_device = string_or_object_state(
    name='jax_default_device',
    default=None,
    help=(
        'Configure the default device for JAX operations. Set to a Device '
        'object (e.g. ``jax.devices("cpu")[0]``) to use that Device as the '
        'default device for JAX operations and jit\'d function calls (there is '
        'no effect on multi-device computations, e.g. pmapped function calls). '
        'Set to None to use the system default device. See '
        ':ref:`faq-data-placement` for more information on device placement.'),
    update_global_hook=_update_default_device_global,
    update_thread_local_hook=_update_default_device_thread_local,
    validator=_validate_default_device)

def _update_disable_jit_global(val):
  jax_jit.global_state().disable_jit = val

def _update_disable_jit_thread_local(val):
  jax_jit.thread_local_state().disable_jit = val

disable_jit = bool_state(
    name='jax_disable_jit',
    default=False,
    help=('Disable JIT compilation and just call original Python.'),
    update_global_hook=_update_disable_jit_global,
    update_thread_local_hook=_update_disable_jit_thread_local)


numpy_rank_promotion = enum_state(
    name='jax_numpy_rank_promotion',
    enum_values=['allow', 'warn', 'raise'],
    default='allow',
    help=('Control NumPy-style automatic rank promotion broadcasting '
          '("allow", "warn", or "raise").'),
    include_in_jit_key=True)

default_matmul_precision = optional_enum_state(
    name='jax_default_matmul_precision',
    enum_values=[
        # Legacy precision API values
        'default', 'high', 'highest', 'bfloat16', 'tensorfloat32', 'float32',
        # Dot algorithm presets
        'ANY_F8_ANY_F8_F32', 'ANY_F8_ANY_F8_F32_FAST_ACCUM', 'ANY_F8_ANY_F8_ANY',
        'ANY_F8_ANY_F8_ANY_FAST_ACCUM', 'F16_F16_F16', 'F16_F16_F32',
        'BF16_BF16_BF16', 'BF16_BF16_F32', 'BF16_BF16_F32_X3',
        'BF16_BF16_F32_X6', 'TF32_TF32_F32', 'TF32_TF32_F32_X3', 'F32_F32_F32',
        'F64_F64_F64',
    ],
    default=None,
    help=('Control the default matmul and conv precision for 32bit inputs.\n\n'

          'Some platforms, like TPU, offer configurable precision levels for '
          'matrix multiplication and convolution computations, trading off '
          'accuracy for speed. The precision can be controlled for each '
          'operation; for example, see the :func:`jax.lax.conv_general_dilated` '
          'and :func:`jax.lax.dot` docstrings. But it can be useful to control '
          'the default behavior obtained when an operation is not given a '
          'specific precision.\n\n'

          'This option can be used to control the default precision '
          'level for computations involved in matrix multiplication and '
          'convolution on 32bit inputs. The levels roughly describe the '
          "precision at which scalar products are computed. The 'bfloat16' "
          "option is the fastest and least precise; 'float32' is similar to "
          "full float32 precision; 'tensorfloat32' is intermediate.\n\n"

          'This parameter can also be used to specify an accumulation '
          '"algorithm" for functions that perform matrix multiplications, like '
          ':func:`jax.lax.dot`. To specify an algorithm, set this option to '
          'the name of a :class:`~jax.lax.DotAlgorithmPreset`.\n\n'),
    include_in_jit_key=True)


traceback_filtering = enum_state(
    name = 'jax_traceback_filtering',
    enum_values=["off", "tracebackhide", "remove_frames", "quiet_remove_frames",
                 "auto"],
    default="auto",
    help="Controls how JAX filters internal frames out of tracebacks.\n\n"
         "Valid values are:\n"
         " * \"off\": disables traceback filtering.\n"
         " * \"auto\": use \"tracebackhide\" if running under a sufficiently"
         " new IPython, or \"remove_frames\" otherwise.\n"
         " * \"tracebackhide\": adds \"__tracebackhide__\" annotations to"
         " hidden stack frames, which some traceback printers support.\n"
         " * \"remove_frames\": removes hidden frames from tracebacks, and adds"
         " the unfiltered traceback as a __cause__ of the exception.\n"
         " * \"quiet_remove_frames\": removes hidden frames from tracebacks, and adds"
         " a brief message (to the __cause__ of the exception) describing that this has"
         " happened.\n")

# This flag is for internal use.
# TODO(tianjianlu): Removes once we always enable cusparse lowering.
# TODO(b/262050896): Set to true after bug is fixed
bcoo_cusparse_lowering = bool_state(
    name='jax_bcoo_cusparse_lowering',
    default=False,
    help=('Enables lowering BCOO ops to cuSparse.'))

# TODO(mattjj): remove this flag when we ensure we only succeed at trace-staging
# if the intended backend can handle lowering the result
dynamic_shapes = bool_state(
    name='jax_dynamic_shapes',
    default=bool(os.getenv('JAX_DYNAMIC_SHAPES', '')),
    help=('Enables experimental features for staging out computations with '
          'dynamic shapes.'),
    include_in_jit_key=True)

# This is for stackless backward compat with e.g. equinox
eager_constant_folding = bool_state(
    name='eager_constant_folding',
    default=False,
    help=('Attempt constant folding during staging.'),
    include_in_jit_key=True)

# This flag is temporary during rollout of the remat barrier.
# TODO(parkers): Remove if there are no complaints.
remat_opt_barrier = bool_state(
    name='jax_remat_opt_barrier',
    default=True,
    help=('Enables using optimization-barrier op for lowering remat.'))

enable_remat_opt_pass = bool_state(
    name='jax_compiler_enable_remat_pass',
    default=True,
    help=('Config to enable / disable the rematerialization HLO pass. '
          'Useful to allow XLA to automatically trade off memory and '
          'compute when encountering OOM errors. However, you are '
          'likely to get better results manually with jax.checkpoint'))

# TODO(sharadmv,mattjj): set default to True, then remove
eager_pmap = bool_state(
    name='jax_eager_pmap',
    default=True,
    upgrade=True,
    help='Enable eager-mode pmap when jax_disable_jit is activated.')

no_tracing = bool_state(
    name='jax_no_tracing',
    default=False,
    help='Disallow tracing for JIT compilation.')

disable_vmap_shmap_error = bool_state(
    name='jax_disable_vmap_shmap_error',
    default=False,
    upgrade=False,
    help='Temporary workaround to disable an error check in vmap-of-shmap.')

# TODO(mattjj): remove once we land mutable array plumbing, or face great shame
custom_vjp_disable_shape_check = bool_state(
    name='jax_custom_vjp_disable_shape_check',
    default=False,
    upgrade=True,
    help='Disable the check from #19009 to enable some custom_vjp hacks.')

mutable_array_checks = bool_state(
    name='jax_mutable_array_checks',
    default=False,
    upgrade=True,
    help='Enable error checks for mutable arrays that rule out aliasing.')

xla_runtime_errors = bool_state(
    name='jax_experimental_unsafe_xla_runtime_errors',
    default=False,
    help=('Enable XLA runtime errors for jax.experimental.checkify.checks '
          'on CPU and GPU. These errors are async, might get lost and are not '
          'very readable. But, they crash the computation and enable you '
          'to write jittable checks without needing to checkify. Does not '
          'work under pmap/pjit.')
)

jax_xla_profile_version = int_state(
    name='jax_xla_profile_version',
    default=0,
    help=(
        'Optional profile version for XLA compilation. This is meaningful '
        'only when XLA is configured to support the remote compilation '
        'profile feature.'),
    include_in_jit_key=True,
)

@contextlib.contextmanager
def explicit_device_put_scope() -> Iterator[None]:
  """Indicates that the current context is an explicit device_put*() call."""
  state = guard_lib.thread_local_state()
  prev = state.explicit_device_put
  state.explicit_device_put = True
  try:
    yield
  finally:
    state.explicit_device_put = prev

@contextlib.contextmanager
def explicit_device_get_scope() -> Iterator[None]:
  """Indicates that the current context is an explicit device_get() call."""
  state = guard_lib.thread_local_state()
  prev = state.explicit_device_get
  state.explicit_device_get = True
  try:
    yield
  finally:
    state.explicit_device_get = prev

def _update_transfer_guard(state, key, val):
  """Applies the transfer guard level within guard_lib."""
  if val is None:
    setattr(state, key, None)
  elif val == 'allow':
    setattr(state, key, guard_lib.TransferGuardLevel.ALLOW)
  elif val == 'log':
    setattr(state, key, guard_lib.TransferGuardLevel.LOG)
  elif val == 'disallow':
    setattr(state, key, guard_lib.TransferGuardLevel.DISALLOW)
  elif val == 'log_explicit':
    setattr(state, key, guard_lib.TransferGuardLevel.LOG_EXPLICIT)
  elif val == 'disallow_explicit':
    setattr(state, key, guard_lib.TransferGuardLevel.DISALLOW_EXPLICIT)
  else:
    assert False, f'Invalid transfer guard level {val}'

transfer_guard_host_to_device = optional_enum_state(
    name='jax_transfer_guard_host_to_device',
    enum_values=[
        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
    ],
    # The default is applied by guard_lib. Use None here to avoid accidentally
    # overriding --jax_transfer_guard.
    default=None,
    help=('Select the transfer guard level for host-to-device transfers. '
          'Default is "allow".'),
    update_global_hook=lambda val: _update_transfer_guard(
        guard_lib.global_state(), 'host_to_device', val),
    update_thread_local_hook=lambda val: _update_transfer_guard(
        guard_lib.thread_local_state(), 'host_to_device', val))

transfer_guard_device_to_device = optional_enum_state(
    name='jax_transfer_guard_device_to_device',
    enum_values=[
        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
    ],
    # The default is applied by guard_lib. Use None here to avoid accidentally
    # overriding --jax_transfer_guard.
    default=None,
    help=('Select the transfer guard level for device-to-device transfers. '
          'Default is "allow".'),
    update_global_hook=lambda val: _update_transfer_guard(
        guard_lib.global_state(), 'device_to_device', val),
    update_thread_local_hook=lambda val: _update_transfer_guard(
        guard_lib.thread_local_state(), 'device_to_device', val))

transfer_guard_device_to_host = optional_enum_state(
    name='jax_transfer_guard_device_to_host',
    enum_values=[
        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
    ],
    # The default is applied by guard_lib. Use None here to avoid
    # accidentally overriding --jax_transfer_guard.
    default=None,
    help=('Select the transfer guard level for device-to-host transfers. '
          'Default is "allow".'),
    update_global_hook=lambda val: _update_transfer_guard(
        guard_lib.global_state(), 'device_to_host', val
    ),
    update_thread_local_hook=lambda val: _update_transfer_guard(
        guard_lib.thread_local_state(), 'device_to_host', val))

def _update_all_transfer_guard_global(val):
  for name in ('jax_transfer_guard_host_to_device',
               'jax_transfer_guard_device_to_device',
               'jax_transfer_guard_device_to_host'):
    config.update(name, val)

_transfer_guard = optional_enum_state(
    name='jax_transfer_guard',
    enum_values=[
        'allow', 'log', 'disallow', 'log_explicit', 'disallow_explicit'
    ],
    # The default is applied by guard_lib. Use None here to avoid accidentally
    # overriding --jax_transfer_guard_*.
    default=None,
    help=('Select the transfer guard level for all transfers. This option is '
          'set-only; the transfer guard level for a specific direction should '
          'be read using the per-transfer direction option. '
          'Default is "allow".'),
    update_global_hook=_update_all_transfer_guard_global)

@contextlib.contextmanager
def transfer_guard(new_val: str) -> Iterator[None]:
  """A contextmanager to control the transfer guard level for all transfers.

  For more information, see
  https://jax.readthedocs.io/en/latest/transfer_guard.html

  Args:
    new_val: The new thread-local transfer guard level for all transfers.

  Yields:
    None.
  """
  with contextlib.ExitStack() as stack:
    stack.enter_context(transfer_guard_host_to_device(new_val))
    stack.enter_context(transfer_guard_device_to_device(new_val))
    stack.enter_context(transfer_guard_device_to_host(new_val))
    stack.enter_context(_transfer_guard(new_val))
    yield


def _update_garbage_collection_guard(state, key, val):
  """Applies the transfer guard level within guard_lib."""
  if val is None:
    setattr(state, key, None)
  elif val == 'allow':
    setattr(state, key, guard_lib.GarbageCollectionGuardLevel.ALLOW)
  elif val == 'log':
    setattr(state, key, guard_lib.GarbageCollectionGuardLevel.LOG)
  elif val == 'fatal':
    setattr(state, key, guard_lib.GarbageCollectionGuardLevel.FATAL)
  else:
    assert False, f'Invalid garbage collection guard level {val}'

array_garbage_collection_guard = optional_enum_state(
    name='jax_array_garbage_collection_guard',
    enum_values=['allow', 'log', 'fatal'],
    # The default is applied by guard_lib.
    default=None,
    help=(
        'Select garbage collection guard level for "jax.Array" objects.\nThis'
        ' option can be used to control what happens when a "jax.Array"'
        ' object is garbage collected. It is desirable for "jax.Array"'
        ' objects to be freed by Python reference couting rather than garbage'
        ' collection in order to avoid device memory being held by the arrays'
        ' until garbage collection occurs.\n\nValid values are:\n * "allow":'
        ' do not log garbage collection of "jax.Array" objects.\n * "log":'
        ' log an error when a "jax.Array" is garbage collected.\n * "fatal":'
        ' fatal error if a "jax.Array" is garbage collected.\nDefault is'
        ' "allow". Note that not all cycles may be detected.'
    ),
    update_global_hook=lambda val: _update_garbage_collection_guard(
        guard_lib.global_state(), 'garbage_collect_array', val
    ),
    update_thread_local_hook=lambda val: _update_garbage_collection_guard(
        guard_lib.thread_local_state(), 'garbage_collect_array', val
    ),
)

# Don't define a context manager since this isn't threadsafe.
string_state(
    name='jax_debug_log_modules',
    default='',
    help=('Comma-separated list of module names (e.g. "jax" or '
          '"jax._src.xla_bridge,jax._src.dispatch") to enable debug logging '
          'for.'),
    update_global_hook=logging_config.update_debug_log_modules)

# Don't define a context manager since this isn't threadsafe.
optional_enum_state(
    name='jax_logging_level',
    enum_values=['NOTSET', 'DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'],
    default=logging.getLevelName(logging.getLogger("jax").level),
    help=('Set the corresponding logging level on all jax loggers. Only string'
          ' values from ["NOTSET", "DEBUG", "INFO", "WARNING", "ERROR",'
          ' "CRITICAL"] are accepted. If None, the logging level will not be'
          ' set. Includes C++ logging.'),
    update_global_hook=lambda logging_level: \
      logging_config.update_logging_level_global(logging_level=logging_level)
)

pmap_no_rank_reduction = bool_state(
    name='jax_pmap_no_rank_reduction',
    default=True,
    help='If True, pmap shards have a the same rank as their enclosing array.',
)

use_shardy_partitioner = bool_state(
    name='jax_use_shardy_partitioner',
    default=False,
    upgrade=True,
    help=(
        'Whether to lower to Shardy. Shardy is a new open sourced propagation '
        'framework for MLIR. Currently Shardy is experimental in JAX. See '
        'www.github.com/openxla/shardy'
    ),
    include_in_jit_key=True,
)

gpu_use_magma = enum_state(
    name='jax_use_magma',
    enum_values=['off', 'on', 'auto'],
    default='auto',
    help=(
        'Enable experimental support for MAGMA-backed lax.linalg.eig on GPU. '
        'See the documentation for lax.linalg.eig for more details about how '
        'to use this feature.'
    ),
)

exec_time_optimization_effort = float_state(
    name='jax_exec_time_optimization_effort',
    default=0.0,
    help='Effort for minimizing execution time (higher means more effort), valid range [-1.0, 1.0].'
)

memory_fitting_effort = float_state(
    name='jax_memory_fitting_effort',
    default=0.0,
    help='Effort for minimizing memory usage (higher means more effort), valid range [-1.0, 1.0].'
)

optimization_level = enum_state(
    name='jax_optimization_level',
    enum_values=[
        'UNKNOWN',
        'O0',
        'O1',
        'O2',
        'O3',
    ],
    default='UNKNOWN',
    help='The degree to which the compiler should optimize for execution time',
    include_in_jit_key=True
)

memory_fitting_level = enum_state(
    name='jax_memory_fitting_level',
    enum_values=[
        'UNKNOWN',
        'O0',
        'O1',
        'O2',
        'O3',
    ],
    default='UNKNOWN',
    help=(
        'The degree to which the compiler should attempt to make the program'
        ' fit in memory'
    ),
    include_in_jit_key=True
)

cpu_collectives_implementation = optional_enum_state(
    name='jax_cpu_collectives_implementation',
    enum_values=["gloo", "mpi", "megascale"],
    default=None,
    help=(
        "Cross-process collective implementation used on CPU. Must be one of "
        '("gloo", "mpi")'),
)

num_cpu_devices = int_state(
    name="jax_num_cpu_devices",
    default=-1,
    help=(
        "Number of CPU devices to use. If not provided, the value of "
        "the XLA flag --xla_force_host_platform_device_count is used."
        " Must be set before JAX is initialized."),
)

enable_empty_arrays = bool_state(
    name='jax_enable_empty_arrays',
    default=False,
    help=(
        "Enable the creation of an Array from an empty list of single-device "
        "arrays. This is to support MPMD/pipeline parallelism in McJAX (WIP)."
    )
)

use_high_dynamic_range_gumbel = bool_state(
    name='jax_high_dynamic_range_gumbel',
    default=False,
    help='If True, gumble noise draws two samples to cover low probability '
         'events with more precision.',
)