{"record":{"id":"401e41189cf9f2d7","repo":"jax-ml/jax","slug":"invalid-flag-value-value","errorCode":null,"errorMessage":"invalid flag value: {value}","messagePattern":"invalid flag value: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_custom_call.py","lineNumber":383,"sourceCode":"      config.write(str(self.vmem_limit_bytes).encode(\"ascii\"))\n      config.write(b'}]')\n    if self.flags is not None:\n      config.write(b', \"flag_configs\": [')\n      for i, (flag, value) in enumerate(self.flags.items()):\n        config.write(b'{\"flag_type\": \"')\n        config.write(flag.encode(\"ascii\"))\n        config.write(b'\", \"value\": {')\n        if isinstance(value, bool):\n          config.write(b'\"boolean_value\": ')\n          config.write(b\"true\" if value else b\"false\")\n        elif isinstance(value, int):\n          config.write(b'\"integer_value\": ')\n          config.write(str(value).encode(\"ascii\"))\n        elif isinstance(value, float):\n          config.write(b'\"double_value\": ')\n          config.write(str(value).encode(\"ascii\"))\n        else:\n          raise ValueError(\"invalid flag value: \" + str(value))\n        config.write(b\"}}\")\n        if i + 1 != len(self.flags):\n          config.write(b\",\")\n      config.write(b\"]\")\n    if self.device_type == \"sparsecore\" and self.active_core_count == 1:\n      config.write(b', \"megachip_parallelism_config\": {\"cores\": [\"0\"]}')\n    config.write(b\"}\")\n    return config.getvalue()\n\n\ndef _compact_json_object(**kwargs: Any) -> bytes:\n  return json.dumps(\n      kwargs, sort_keys=True, indent=0, separators=(\",\", \":\")\n  ).encode(\"ascii\")\n\n\n@tpu_custom_call_p.def_abstract_eval\ndef _tpu_custom_call_abstract_eval(*_, out_avals, **__):","sourceCodeStart":365,"sourceCodeEnd":401,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_custom_call.py#L365-L401","documentation":"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.","triggerScenarios":"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.","commonSituations":"Passing arbitrary Python objects as custom_call flags; API drift where flags were previously ignored.","solutions":["Inspect the flags dict passed to the custom call and convert each value to bool/int/float","Update JAX/Pallas to a version matching the kernel's flag API"],"exampleFix":"// before\nflags = {\"layers\": [1, 2]}\n// after\nflags = {\"num_layers\": 2}","handlingStrategy":"validation","validationCode":"flags = {...}\nassert all(isinstance(v, (bool, int, float)) and not isinstance(v, bool) or isinstance(v, bool) for v in flags.values())","typeGuard":"def is_valid_flag_value(v) -> bool:\n    return isinstance(v, (bool, int, float))","tryCatchPattern":null,"preventionTips":["Keep custom-call flags to bool/int/float only","Add a unit test asserting your kernel's flags serialize"],"tags":["tpu","pallas","custom-call","config-serialization","jax"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}