jax-ml/jax · error · ValueError

new enum value must be None or in {enum_values}, got {new_va

Error message

new enum value must be None or in {enum_values}, got {new_val} of type {type(new_val)}.

What it means

The parser for optional_enum_state rejects new values that are neither None nor a str in enum_values. Raised when entering the context manager or calling config.update for such a flag.

Source

Thrown at jax/_src/config.py:579

    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 parser(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)}.")
    return new_val

  s = State[str | None](
      name, default, help, update_global_hook, update_thread_local_hook,
      parser, include_in_jit_key=include_in_jit_key,
      include_in_trace_context=include_in_trace_context,
  )
  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 enum_class_state[EnumType: enum.Enum](

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass None or the exact string in enum_values (use MyEnum.FOO.value)
Validate/whitelist user-supplied values against the flag's enum_values before applying

Example fix

// before
with jax.my_flag(user_choice):  # user_choice = 'auto'
    ...
// after
if user_choice not in ('foo', 'bar', None):
    raise ValueError(f'bad choice: {user_choice}')
with jax.my_flag(user_choice):
    ...

Defensive patterns

Strategy: validation

Validate before calling

if new_val is not None and (type(new_val) is not str or new_val not in enum_values):
    raise ValueError(f'expected None or one of {enum_values}, got {new_val!r}')

Type guard

def is_optional_enum_value(v, values) -> bool:
    return v is None or (type(v) is str and v in values)

Try / catch

try:
    with my_flag(v):
        ...
except ValueError:
    with my_flag(None):
        ...

Prevention

Centralize flag-setting behind a helper that validates against enum_values
Never pass raw user input directly to config context managers

When it happens

Trigger: with my_flag(MyEnum.FOO): or jax.config.update('jax_my_flag', 'typo') where 'typo' is not an allowed value.

Common situations: Dynamically choosing a config value from user input or an enum object and passing it straight to an optional enum flag.

Understand the failure class

Background: "Invalid value" and "allowed values are" config errors: what your library rejected and how to fix it — this error's family across 41 libraries.

Related errors

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