jax-ml/jax · error · TypeError

Default value must be of type bool, got {default} of type {g

Error message

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

What it means

bool_state, the helper that defines boolean JAX config options, validates that its default is a Python bool before reading the environment variable override. Because the option's semantic domain is strictly boolean, passing an int, string, or None default is rejected at definition time with a message showing the offending value and its type.

Source

Thrown at jax/_src/config.py:408

        name='jax_enable_foo',
        default=False,
        help='Enable foo.')

    # Now the JAX_ENABLE_FOO shell environment variable and --jax_enable_foo
    # command-line flag can be used to control the process-level value of
    # the configuration option, in addition to using e.g.
    # ``config.update("jax_enable_foo", True)`` directly. We can also use a
    # context manager:

    with enable_foo(True):
      ...

  The value of the thread-local state or flag can be accessed via
  ``config.jax_enable_foo``. Reading it via ``config.FLAGS.jax_enable_foo`` is
  an error.
  """
  if not isinstance(default, bool):
    raise TypeError(f"Default value must be of type bool, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  default = bool_env(name.upper(), default)
  name = name.lower()
  if upgrade:
    help += ' ' + UPGRADE_BOOL_HELP
    extra_description += UPGRADE_BOOL_EXTRA_DESC
  config._contextmanager_flags.add(name)

  def parser(val):
    if validator:
      validator(val)
    return bool(val)

  s = State[bool](
      name, default, help, update_global_hook=update_global_hook,
      update_thread_local_hook=update_thread_local_hook,
      extra_description=extra_description, default_context_manager_value=True,
      parser=parser, include_in_jit_key=include_in_jit_key,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Pass a literal bool: default=False or default=True
  2. If the default is computed, coerce: bool(int(os.getenv('MY_FLAG', '0')))
  3. For string-parsed env defaults, let bool_env handle it by just supplying the bool default

Example fix

# before
jax.config.define_bool_state('jax_my_flag', int(os.getenv('MY_FLAG', '0')), 'help')
# after
jax.config.define_bool_state('jax_my_flag', os.getenv('MY_FLAG', '0') == '1', 'help')
Defensive patterns

Strategy: type-guard

Validate before calling

default = os.getenv('MY_FLAG', '0') == '1'  # ensure bool
assert isinstance(default, bool), f'default must be bool, got {type(default).__name__}'
jax.config.define_bool_state('jax_my_flag', default, 'help')

Type guard

def is_bool(v: object) -> TypeGuard[bool]:
    return isinstance(v, bool)

Prevention

When it happens

Trigger: Calling jax.config.define_bool_state / bool_state('jax_my_flag', 1, ...), default='true', or default=None. Also happens when a default is computed from env/config code that may return non-bool.

Common situations: Writing custom flags in plugins or forks; porting code where the default came from an env read (os.getenv returns a string); using 0/1 as 'boolean' defaults out of habit.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors


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