jax-ml/jax · error · ValueError

new float config value must be None or of type float, got {n

Error message

new float config value must be None or of type float, got {new_val} of type {type(new_val)}

What it means

The parser for float_state accepts None, float, or int new values; anything else (str, bool-as-intended, objects) raises this ValueError when the flag is set.

Source

Thrown at jax/_src/config.py:768

  Returns:
    A contextmanager to control the thread-local state value.
  """
  if not isinstance(default, float):
    raise TypeError(f"Default value must be of type float, got {default} "
                    f"of type {getattr(type(default), '__name__', type(default))}")
  name = name.lower()
  default_env = os.getenv(name.upper())
  if default_env is not None:
    try:
      default = float(default_env)
    except ValueError:
      raise ValueError(f"Invalid value \"{default_env}\" for JAX flag {name}")
  config._contextmanager_flags.add(name)

  def parser(new_val):
    if new_val is not None and not isinstance(new_val, (float, int)):
      raise ValueError(
        f'new float config value must be None or of type float, '
        f'got {new_val} of type {type(new_val)}')
    return new_val

  s = State[float](name, default, help, update_global_hook,
                   update_thread_local_hook, parser)
  config.add_option(name, s, float, meta_args=[], meta_kwargs={"help": help})
  setattr(Config, name, property(lambda _: s.value))
  return s


def string_state(
    name: str,
    default: str,
    help: str,
    *,
    update_global_hook: Callable[[str], None] | None = None,
    update_thread_local_hook: Callable[[str | None], None] | None = None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Convert with float(value) before setting
  2. Use argparse type=float or a schema validator for config files

Example fix

// before
jax.config.update('jax_my_flag', cfg['rate'])  # '0.5' from JSON
// after
jax.config.update('jax_my_flag', float(cfg['rate']))
Defensive patterns

Strategy: validation

Validate before calling

if new_val is not None and not isinstance(new_val, (float, int)):
    raise ValueError(f'expected float/int or None, got {type(new_val)}')

Type guard

def is_float_like(v) -> bool:
    return v is None or isinstance(v, (float, int)) and not isinstance(v, bool)

Try / catch

try:
    jax.config.update('jax_my_flag', val)
except ValueError:
    jax.config.update('jax_my_flag', float(val))

Prevention

When it happens

Trigger: jax.config.update('jax_my_flag', '0.5') — a string instead of a number; or passing a numpy scalar of dtype object.

Common situations: Config-driven scripts passing values straight from JSON/YAML strings or argparse without float conversion.

Understand the failure class

Background: "Invalid value" and "allowed values are" config errors: what your library rejected and how to fix it — this error's family across 41 libraries.

Related errors


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