jax-ml/jax · error · ValueError

Unsupported core type: {core_type}

Error message

Unsupported core type: {core_type}

What it means

memory_space_to_tpu_memory_space maps a Pallas memory space to a TPU-specific space based on the core type. Only TensorCore and SparseCore vector subcore map to VMEM and SparseCore scalar subcore maps to SMEM; any other CoreType value falls through to this ValueError.

Source

Thrown at jax/_src/pallas/mosaic/core.py:499

    | pallas_core.MemorySpace
    | pallas_core.CoreMemorySpace
    | jax_core.MemorySpace
):
  match memory_space:
    case None:
      match core_type:
        case CoreType.TC:
          return pallas_core.MemorySpace.ANY
        case CoreType.SC_SCALAR_SUBCORE | CoreType.SC_VECTOR_SUBCORE:
          return MemorySpace.HBM
    case pallas_core.MemorySpace.DEFAULT:
      match core_type:
        case CoreType.TC | CoreType.SC_VECTOR_SUBCORE:
          return MemorySpace.VMEM
        case CoreType.SC_SCALAR_SUBCORE:
          return MemorySpace.SMEM
        case _:
          raise ValueError(f"Unsupported core type: {core_type}")
    case pallas_core.MemorySpace.ANY | jax_core.MemorySpace.Host:
      return memory_space
    case (
        pallas_core.MemorySpace.ERROR
        | pallas_core.MemorySpace.INDEX
        | pallas_core.MemorySpace.KEY
    ):
      return MemorySpace.SMEM
    case pallas_core.CoreMemorySpace():
      return (
          memory_space.memory_space
          if memory_space.mesh.core_type is core_type
          else memory_space
      )
    case acc if isinstance(acc, AccMemorySpace):
      return acc
    case MemorySpace():
      return memory_space

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Update JAX (and its bundled Mosaic) to matching versions so all CoreType members are handled
2. Avoid annotating refs with core types outside TC / SC vector / SC scalar subcores on TPU
3. File an issue if a legitimate new core type is unhandled in this match statement

from jax._src.pallas.mosaic.core import CoreType
assert core_type in (CoreType.TC, CoreType.SC_VECTOR_SUBCORE, CoreType.SC_SCALAR_SUBCORE)

When it happens

Trigger: Lowering a Mosaic kernel (pipelined/unpipelined lowering, block-mapping checks, load lowering) where a memory space carries a CoreType not in {TC, SC_VECTOR_SUBCORE, SC_SCALAR_SUBCORE}, e.g., a newly added or synthetic core type.

Common situations: Version skew between JAX and a newer/older Mosaic/PJRT that introduces new CoreType enum members; custom memory-space annotations with unexpected core types.

Related errors

• Memory space {self.memory_space} is not supported by mesh {s
• 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.
• Loads and stores are only allowed on VMEM and SMEM reference

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