jax-ml/jax · error · NotImplementedError

input_memory_space_colors only supports HBM, VMEM and SMEM

Error message

input_memory_space_colors only supports HBM, VMEM and SMEM

What it means

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.

Source

Thrown at jax/_src/tpu_custom_call.py:321

          config.write(f'{{"color":{memory_space.color}}}'.encode("ascii"))
        comma = True
      if comma:
        config.write(b"]")
    if self.input_memory_spaces is not None:
      comma = False
      for i, memory_space in enumerate(self.input_memory_spaces):
        if memory_space is None:
          continue
        if memory_space is MemorySpace.SMEM:
          # TODO(sharadmv): Add support for SMEM (though atm, XLA will not
          # page out SMEM arrays).
          continue
        if memory_space not in (
            MemorySpace.HBM,
            MemorySpace.VMEM,
            MemorySpace.SMEM,
        ):
          raise NotImplementedError(
              "input_memory_space_colors only supports HBM, VMEM and SMEM"
          )
        if comma:
          config.write(b",")
        else:
          config.write(b', "input_memory_space_colors": [')
        config.write(
            f'{{"operand_index":{i},"color":{memory_space.color}}}'
            .encode("ascii")
        )
        comma = True
      if comma:
        config.write(b"]")
    if self.disable_bounds_checks:
      config.write(b', "disable_bounds_checks": ')
      config.write(str(self.disable_bounds_checks).lower().encode("ascii"))
    if self.disable_semaphore_checks:
      config.write(b', "disable_semaphore_checks": ')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restrict input operand memory spaces to MemorySpace.HBM, VMEM, or SMEM.
2. Remove the custom input memory space specification and use defaults.
3. Update JAX to a version where the serializer supports your memory space, or patch upstream.

ALLOWED = {MemorySpace.HBM, MemorySpace.VMEM, MemorySpace.SMEM}
assert all(sp in ALLOWED for sp in input_memory_spaces), 'HBM/VMEM/SMEM only'

def supported_input_space(sp) -> bool:
    return sp in (MemorySpace.HBM, MemorySpace.VMEM, MemorySpace.SMEM)

try:
    lower/custom_call(...)
except NotImplementedError as e:
    if 'input_memory_space_colors' in str(e): retry without custom memory spaces
    else: raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• tpu_custom_call does not support non-trivial batching.
• QDWH implementation is only supported on TPU
• Failed to find assignment for logical_axis_index {logical_ax
• masked load_p
• run_scoped_p with collective axes is not supported

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