jax-ml/jax · error · TypeError

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

Error message

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

What it means

optional_enum_state requires `default` to be a str or None. Any other type (Enum member, int, bool) raises this TypeError at flag-definition time.

Source

Thrown at jax/_src/config.py:568

  """Set up thread-local state and return a contextmanager for managing it.

  See docstring for ``bool_state``.

  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: optional string, default value.
    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,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass MyEnum.FOO.value or None
Or use enum_class_state if you want to work with Enum members

Example fix

// before
config.optional_enum_state('my_flag', default=MyEnum.FOO, enum_values=['foo','bar'])
// after
config.optional_enum_state('my_flag', default=None, enum_values=['foo','bar'])

Defensive patterns

Strategy: type-guard

Validate before calling

assert default is None or isinstance(default, str)

Type guard

def is_optional_str(v) -> bool:
    return v is None or isinstance(v, str)

Prevention

Convert enum members with .value before using as defaults
Use enum_class_state when you want typed Enum values

When it happens

Trigger: Calling config.optional_enum_state('my_flag', default=MyEnum.FOO, enum_values=[...]) instead of default='foo' or default=None.

Common situations: Defining an optional enum flag (e.g. a 'not set means auto' option) and passing an enum object rather than its string value.

Related errors

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