jax-ml/jax · error · TypeError

Context manager for {state.name} config option requires

Error message

Context manager for {state.__name__} config option requires an argument representing the new value for the config option.

What it means

Context managers generated for config options require the new value as an argument (e.g. jax.config.enable_x64(True) style, or with_x64(True)) unless the option was defined with a default_context_manager_value. If neither an argument nor a constructor default exists, __init__ raises TypeError immediately rather than silently using the current value.

Source

Thrown at jax/_src/config.py:301

    Used to avoid cyclic import dependencies."""
    self._update_thread_local_hook = update_thread_local_hook
    self._update_global_hook = update_global_hook
    update_global_hook(self.get_global())


class StateContextManager[FuncType: Callable[..., Any]]:
  __slots__ = ['state', 'new_val', 'prev']

  def __init__(self, state, new_val):
    self.state = state

    if new_val is no_default:
      if state._default_context_manager_value is not no_default:
        new_val = state._default_context_manager_value  # default_context_manager_value provided to constructor
      else:
        # no default_value provided to constructor and no value provided as an
        # argument, so we raise an error
        raise TypeError(f"Context manager for {state.__name__} config option "
                        "requires an argument representing the new value for "
                        "the config option.")
    if state._parser:
      self.new_val = state._parser(new_val)
    else:
      self.new_val = new_val

  def __enter__(self):
    self.prev = self.state.swap_local(self.new_val)
    if self.state._update_thread_local_hook:
      self.state._update_thread_local_hook(self.new_val)

  def __exit__(self, exc_type, exc_value, traceback):
    self.state.set_local(self.prev)
    if self.state._update_thread_local_hook:
      if self.prev is config_ext.unset:
        self.state._update_thread_local_hook(None)
      else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass the new value explicitly: jax.config.with_x64(True)
Check the option definition: only options with default_context_manager_value can be called with no args
If you own the option, provide default_context_manager_value when defining it

Example fix

# before
with jax.config.some_ctx():  # TypeError
    ...
# after
with jax.config.some_ctx(True):
    ...

Defensive patterns

Strategy: validation

Validate before calling

import inspect

def with_cfg(ctx_mgr, *args, **kwargs):
    if not args and not kwargs:
        raise TypeError(f'{ctx_mgr} requires an explicit value, e.g. {ctx_mgr}(True)')
    return ctx_mgr(*args, **kwargs)

with with_cfg(jax.config.some_ctx, True):
    ...

Try / catch

try:
    cm = jax.config.some_ctx()
except TypeError as e:
    if 'requires an argument' in str(e):
        cm = jax.config.some_ctx(True)
    else:
        raise

Prevention

Always pass the target value to config context managers
Check the define_* call for default_context_manager_value to know if argumentless use is allowed
Standardize on explicit values in with-statements for clarity

When it happens

Trigger: Calling a config context manager with no argument: jax.config.some_context_manager() where some_context_manager was defined without default_context_manager_value. Options with defaults (like jax.debug_nans()) work argumentless; ones without do not.

Common situations: Assuming all config context managers are zero-argument toggles; calling with_x64() expecting it to default to True.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

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