jax-ml/jax · error · ValueError

invalid memory space: {self}

Error message

invalid memory space: {self}

What it means

MemorySpace.to_color raises ValueError when asked to convert a MemorySpace enum member that has no assigned color code — i.e. an invalid or newly-added/unknown memory space value falls through the if/elif chain.

Source

Thrown at jax/_src/tpu_custom_call.py:140

      return 0
    elif self == MemorySpace.VMEM:
      return 1
    elif self == MemorySpace.SEMAPHORE_MEM:
      return 2
    elif self == MemorySpace.SC_SCALAR_SEMAPHORE_MEM:
      return 8
    elif self == MemorySpace.SC_VECTOR_SEMAPHORE_MEM:
      return 10
    elif self == MemorySpace.SMEM:
      return 4
    elif self == MemorySpace.HOST:
      return 5
    elif self == MemorySpace.SC_SCALAR_SMEM:
      return 11
    elif self == MemorySpace.SC_VECTOR_SMEM:
      return 12
    else:
      raise ValueError("invalid memory space: " + str(self))


class CostEstimate(TypedDict):
  flops: int
  transcendentals: int
  bytes_accessed: int
  remote_bytes_transferred: int

  def to_json(self) -> bytes:  # pyrefly: ignore[bad-class-definition]
    return (
        f'{{"flops": {self["flops"]}, "transcendentals":'
        f' {self["transcendentals"]}, "bytes_accessed":'
        f' {self["bytes_accessed"]}, "remote_bytes_transferred":'
        f' {self["remote_bytes_transferred"]}}}'
    ).encode("ascii")


class TpuSideEffectType(enum.Enum):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use only documented MemorySpace enum members (HBM, VMEM, SMEM, ...) instead of raw ints.
2. Align JAX versions across processes/components that exchange these values.
3. Catch ValueError and reject/report the unknown space early rather than letting it reach to_color.

VALID = {MemorySpace.HBM, MemorySpace.VMEM, MemorySpace.SMEM}
if space not in VALID: raise ValueError(f'unsupported memory space {space!r}')

def is_valid_memory_space(s) -> bool:
    return isinstance(s, MemorySpace) and s in VALID_SPACES

try:
    color = space.to_color()
except ValueError:
    raise ValueError(f'unknown MemorySpace {space!r}; update JAX or use a documented member')

When it happens

Trigger: Constructing MemorySpace with an out-of-range int or comparing to an unknown member then calling to_color(); typically reached via serialization of layout/cost-analysis attributes that call this conversion.

Common situations: Version mismatch where an enum value exists on one side but not the other; typos in raw integer memory-space codes passed by user code or stale serialized configs.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• Memory space {self.memory_space} is not supported by mesh {s
• Unsupported core type: {core_type}
• Invalid memory space: {memory_space!r}
• The Pallas TPU lowering currently supports in memory space A
• Loads are only allowed on VMEM and SMEM references.

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