jax-ml/jax · error · TypeError

Default value must be of type {enum_class}, got {default} of

Error message

Default value must be of type {enum_class}, got {default} of type {getattr(type(default), '__name__', type(default))}

What it means

enum_class_state requires `default` to be an instance of the given enum_class (not a raw string). Passing a string or other object raises this TypeError at flag definition.

Source

Thrown at jax/_src/config.py:632

    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_class: a subtype of enum.Enum.
    default: an instance of enum_class that is the 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.
    include_in_trace_context: bool, optional: whether to include the state in
      the trace context.
    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, enum_class):
    raise TypeError(
        f'Default value must be of type {enum_class}, got {default} '
        f"of type {getattr(type(default), '__name__', type(default))}"
    )
  name = name.lower()
  default_str = os.getenv(name.upper(), None)
  if default_str is not None:
    try:
      default = enum_class(default_str)
    except ValueError as e:
      raise ValueError(f"Invalid value \"{default_str}\" for JAX flag {name}") from e
  config._contextmanager_flags.add(name)

  def parser(new_val):
    if isinstance(new_val, str):
      return enum_class(new_val)
    if not isinstance(new_val, enum_class):
      raise TypeError(
          f'new enum value must be an instance of {enum_class}, got'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass an Enum member: default=MyEnum.FOO
2. For string defaults, keep using enum_state instead

Example fix

// before
config.enum_class_state('my_flag', enum_class=MyEnum, default='foo')
// after
config.enum_class_state('my_flag', enum_class=MyEnum, default=MyEnum.FOO)

assert isinstance(default, enum_class)

import enum
def is_enum_member(v, cls) -> bool:
    return isinstance(v, cls) and type(v) is not str

When it happens

Trigger: Calling config.enum_class_state('my_flag', enum_class=MyEnum, default='foo') instead of default=MyEnum.FOO.

Common situations: Migrating a flag from enum_state (string-based) to enum_class_state and forgetting to convert the default to an Enum member.

Related errors

• Default value must be of type str, got {default} of type {ge
• Default value must be of type str or None, got {default} of
• new enum value must be an instance of {enum_class}, got {new
• bool() not supported for instances of type '{0}' (did you me
• Default value must be of type bool, got {default} of type {g

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