jax-ml/jax · error · ValueError

Invalid value "{default}" for JAX flag {name}

Error message

Invalid value "{default}" for JAX flag {name}

What it means

enum_state validates that the flag's default value (after checking the uppercase environment variable override) is one of the declared enum_values. If the env var JAX_<FLAG> is set to something outside enum_values, this ValueError is raised at flag definition (typically at jax import).

Source

Thrown at jax/_src/config.py:509

      option.
    default: string, default value.
    help: string, used to populate the flag help information as well as the
      docstring of the returned context manager.
    include_in_jit_key: bool, optional: whether to include the state in the
      JIT cache key.
    extra_validator: optional function to validate the value of the config
      option.

  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 parser(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)}.")
    if extra_validator is not None:
      extra_validator(new_val)
    return new_val

  s = State[str](
      name,
      default,
      help,
      update_global_hook=update_global_hook,
      update_thread_local_hook=update_thread_local_hook,
      parser=parser,
      include_in_jit_key=include_in_jit_key,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the exact flag name and fix/unset the offending JAX_<FLAG> env var: env | grep JAX_
2. Make sure `default` itself is one of enum_values
3. Update the value to one supported by your jax version (see the flag's help text)

Example fix

# before
export JAX_MY_FLAG=baz  # invalid
# after
export JAX_MY_FLAG=foo  # one of enum_values

import os
val = os.getenv('JAX_MY_FLAG')
if val is not None and val not in enum_values:
    raise RuntimeError(f'JAX_MY_FLAG={val!r} invalid; allowed: {enum_values}')

When it happens

Trigger: Defining config.enum_state('my_flag', default='foo', enum_values=['foo','bar']) while JAX_MY_FLAG=baz is exported in the environment; or passing a default not in enum_values.

Common situations: A stale or misspelled JAX_* environment variable from a previous experiment or an old jax version whose accepted values changed (e.g. JAX_PYTHON_PREALLOCATE spellings, backend flag values).

Related errors

• Invalid value "{default_str}" for JAX flag {name}
• new enum value must be in {enum_values}, got {new_val} of ty
• new enum value must be None or in {enum_values}, got {new_va
• Invalid value "{default_env}" for JAX flag {name}
• invalid truth value {val!r} for environment {varname!r}

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