jax-ml/jax · error · ValueError

invalid flag value: {value}

Error message

invalid flag value: {value}

What it means

Raised when serializing Mosaic/Pallas custom-call backend config flags: a flag value passed to a TPU custom call is neither bool, int, nor float, so it cannot be written as JSON 'flag_value' style.

Source

Thrown at jax/_src/tpu_custom_call.py:383

      config.write(str(self.vmem_limit_bytes).encode("ascii"))
      config.write(b'}]')
    if self.flags is not None:
      config.write(b', "flag_configs": [')
      for i, (flag, value) in enumerate(self.flags.items()):
        config.write(b'{"flag_type": "')
        config.write(flag.encode("ascii"))
        config.write(b'", "value": {')
        if isinstance(value, bool):
          config.write(b'"boolean_value": ')
          config.write(b"true" if value else b"false")
        elif isinstance(value, int):
          config.write(b'"integer_value": ')
          config.write(str(value).encode("ascii"))
        elif isinstance(value, float):
          config.write(b'"double_value": ')
          config.write(str(value).encode("ascii"))
        else:
          raise ValueError("invalid flag value: " + str(value))
        config.write(b"}}")
        if i + 1 != len(self.flags):
          config.write(b",")
      config.write(b"]")
    if self.device_type == "sparsecore" and self.active_core_count == 1:
      config.write(b', "megachip_parallelism_config": {"cores": ["0"]}')
    config.write(b"}")
    return config.getvalue()


def _compact_json_object(**kwargs: Any) -> bytes:
  return json.dumps(
      kwargs, sort_keys=True, indent=0, separators=(",", ":")
  ).encode("ascii")


@tpu_custom_call_p.def_abstract_eval
def _tpu_custom_call_abstract_eval(*_, out_avals, **__):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Inspect the flags dict passed to the custom call and convert each value to bool/int/float
2. Update JAX/Pallas to a version matching the kernel's flag API

Example fix

// before
flags = {"layers": [1, 2]}
// after
flags = {"num_layers": 2}

flags = {...}
assert all(isinstance(v, (bool, int, float)) and not isinstance(v, bool) or isinstance(v, bool) for v in flags.values())

def is_valid_flag_value(v) -> bool:
    return isinstance(v, (bool, int, float))

When it happens

Trigger: Calling a Pallas TPU kernel with backend flags containing a non-scalar value (e.g. a string, list, or None) that reaches CustomCallBackendConfig.to_json.

Common situations: Passing arbitrary Python objects as custom_call flags; API drift where flags were previously ignored.

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

• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}
• Acc ref dtype must be float32 or int32, got {dtype}
• Accumulators are not available on TPU {info.chip_version}

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