jax-ml/jax · error · AttributeError

For flags with a corresponding contextmanager, read their va

Error message

For flags with a corresponding contextmanager, read their value via e.g. `config.{name}` rather than `config.FLAGS.{name}`.

What it means

Flags defined with a corresponding context manager must be read through their property on the config object (e.g. jax.config.jax_enable_x64) rather than through the absl-style FLAGS proxy (config.FLAGS.jax_enable_x64). Reading via FLAGS would bypass thread-local/context-local semantics, so config.read explicitly raises AttributeError to redirect the caller.

Source

Thrown at jax/_src/config.py:103

      ...

    def __setattr__(self, name: str, value: Any) -> None:
      ...

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Replace config.FLAGS.name with config.name (e.g. jax.config.jax_enable_x64)
  2. For programmatic access use jax.config.read(name) only for non-contextmanager flags, or getattr(jax.config, name)
  3. If you need the absl flag value, read the absl flag itself (absl.flags.FLAGS['jax_enable_x64'].value) rather than config.FLAGS

Example fix

# before
if jax.config.FLAGS.jax_enable_x64: ...
# after
if jax.config.jax_enable_x64: ...
Defensive patterns

Strategy: validation

Validate before calling

import jax

def cfg_get(name):
    # always read via attribute, never FLAGS
    return getattr(jax.config, name)

assert cfg_get('jax_enable_x64') in (True, False)

Prevention

When it happens

Trigger: Accessing jax.config.FLAGS.<option_name> for any option registered with a context manager (most modern thread-local JAX options like jax_enable_x64, jax_numpy_rank_promotion).

Common situations: Migrating old code that read config.FLAGS in the pre-thread-local JAX config era; absl flag introspection code iterating over FLAGS attributes.

Related errors


AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/398a0e94a534f633. Report an issue: GitHub.