jax-ml/jax · error · TypeError

Default value must be of type str, got {default} of type {ge

Error message

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

What it means

jax._src.config.enum_state requires the `default` for an enum-backed config flag to be a Python str. Any non-str default (int, Enum member, None) raises this TypeError when the flag is defined.

Source

Thrown at jax/_src/config.py:504

  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.
    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass the string value: default=MyEnum.FOO.value
  2. If you have an Enum class, prefer config.enum_class_state which accepts enum members directly

Example fix

// before
config.enum_state('my_flag', default=MyEnum.FOO, enum_values=['foo','bar'])
// after
config.enum_state('my_flag', default='foo', enum_values=['foo','bar'])
Defensive patterns

Strategy: type-guard

Validate before calling

assert isinstance(default, str) and default in enum_values

Type guard

def is_valid_enum_default(v, values) -> bool:
    return isinstance(v, str) and v in values

Prevention

When it happens

Trigger: Calling config.enum_state('my_flag', default=MyEnum.FOO, enum_values=['foo','bar']) — Enum members must be converted with .value; or passing an int/None default.

Common situations: Defining a new JAX config flag whose options are enumerated strings but reusing an enum object or numeric default from other code.

Related errors


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