{"record":{"id":"64057b0ce4ca5b63","repo":"jax-ml/jax","slug":"input-memory-space-colors-only-supports-hbm-vmem","errorCode":null,"errorMessage":"input_memory_space_colors only supports HBM, VMEM and SMEM","messagePattern":"input_memory_space_colors only supports HBM, VMEM and SMEM","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_custom_call.py","lineNumber":321,"sourceCode":"          config.write(f'{{\"color\":{memory_space.color}}}'.encode(\"ascii\"))\n        comma = True\n      if comma:\n        config.write(b\"]\")\n    if self.input_memory_spaces is not None:\n      comma = False\n      for i, memory_space in enumerate(self.input_memory_spaces):\n        if memory_space is None:\n          continue\n        if memory_space is MemorySpace.SMEM:\n          # TODO(sharadmv): Add support for SMEM (though atm, XLA will not\n          # page out SMEM arrays).\n          continue\n        if memory_space not in (\n            MemorySpace.HBM,\n            MemorySpace.VMEM,\n            MemorySpace.SMEM,\n        ):\n          raise NotImplementedError(\n              \"input_memory_space_colors only supports HBM, VMEM and SMEM\"\n          )\n        if comma:\n          config.write(b\",\")\n        else:\n          config.write(b', \"input_memory_space_colors\": [')\n        config.write(\n            f'{{\"operand_index\":{i},\"color\":{memory_space.color}}}'\n            .encode(\"ascii\")\n        )\n        comma = True\n      if comma:\n        config.write(b\"]\")\n    if self.disable_bounds_checks:\n      config.write(b', \"disable_bounds_checks\": ')\n      config.write(str(self.disable_bounds_checks).lower().encode(\"ascii\"))\n    if self.disable_semaphore_checks:\n      config.write(b', \"disable_semaphore_checks\": ')","sourceCodeStart":303,"sourceCodeEnd":339,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_custom_call.py#L303-L339","documentation":"When serializing a tpu_custom_call's cost-analysis/config to JSON (to_json, invoked from the TPU custom call lowering), input_memory_space_colors only supports HBM, VMEM and SMEM. Specifying an input memory space outside that set (e.g. SC scalar/vector SMEM variants or other enum members) raises NotImplementedError.","triggerScenarios":"Passing tpu_custom_call(..., input_output_aliases or cost_analysis with input memory spaces) where an operand's memory space is something like MemorySpace.SC_SCALAR_SMEM or SC_VECTOR_SMEM, triggering the JSON writer during lowering.","commonSituations":"Advanced Pallas/custom-call tuning that specifies exotic operand placements; upgrading JAX where new memory spaces were added to the enum but not to this serializer.","solutions":["Restrict input operand memory spaces to MemorySpace.HBM, VMEM, or SMEM.","Remove the custom input memory space specification and use defaults.","Update JAX to a version where the serializer supports your memory space, or patch upstream."],"exampleFix":null,"handlingStrategy":"validation","validationCode":"ALLOWED = {MemorySpace.HBM, MemorySpace.VMEM, MemorySpace.SMEM}\nassert all(sp in ALLOWED for sp in input_memory_spaces), 'HBM/VMEM/SMEM only'","typeGuard":"def supported_input_space(sp) -> bool:\n    return sp in (MemorySpace.HBM, MemorySpace.VMEM, MemorySpace.SMEM)","tryCatchPattern":"try:\n    lower/custom_call(...)\nexcept NotImplementedError as e:\n    if 'input_memory_space_colors' in str(e): retry without custom memory spaces\n    else: raise","preventionTips":["Stick to HBM/VMEM/SMEM for custom-call operand placement on TPU.","Check release notes when upgrading JAX if you use exotic memory spaces."],"tags":["jax","tpu","custom-call","serialization","not-implemented"],"backgroundTag":"unsupported-configuration-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}