jax-ml/jax · error · ValueError

Unknown core type: {core_type}

Error message

Unknown core type: {core_type}

What it means

While scanning the lowered Mosaic module, an unknown #tpu.core_type attribute was found — only <tc> (tensor core) and the sparsecore variant are recognized.

Source

Thrown at jax/_src/tpu_custom_call.py:546

  def assign_device_type_based_on_core_type(op: ir.Operation) -> ir.WalkResult:
    nonlocal sparsecore_func_found
    nonlocal tensorcore_func_found
    if op.name == "func.func":
      if "tpu.core_type" in op.attributes:
        core_type = op.attributes["tpu.core_type"]
        if str(core_type) in [
            f"#tpu.core_type<{c}>"
            for c in ["sc_scalar_subcore", "sc_vector_subcore"]
        ]:
          sparsecore_func_found = True
          if tensorcore_func_found:
            return ir.WalkResult.INTERRUPT
          return ir.WalkResult.SKIP
        if str(core_type) == "#tpu.core_type<tc>":
          tensorcore_func_found = True
          return ir.WalkResult.SKIP
        raise ValueError(f"Unknown core type: {core_type}")
    return ir.WalkResult.ADVANCE

  module.operation.walk(
      assign_device_type_based_on_core_type, walk_order=ir.WalkOrder.PRE_ORDER
  )
  if tensorcore_func_found and sparsecore_func_found:
    raise ValueError(
        "A single Mosaic kernel cannot contain both TensorCore and SparseCore"
        " functions."
    )
  if sparsecore_func_found:
    return "sparsecore"
  return None


def _get_active_core_count(module: ir.Module) -> int | None:

  def get_core_parallel_dim_size(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Align JAX, Pallas, and Mosaic/jaxlib versions
2. Re-lower the kernel with the toolchain matching your JAX version

import jax, jaxlib
assert version_parse(jax.__version__) >= version_parse(min_required)

When it happens

Trigger: A lowered Pallas module containing a core_type attribute other than the supported tc/sparsecore forms, often from a version mismatch between Mosaic/Pallas IR and this JAX.

Common situations: Mixing JAX version with a newer/older Pallas or Mosaic that emits new core types (e.g. new chip types).

Related errors

• Pallas TPU requires a recent libtpu version (at least 0.0.46
• The iteration bounds and dimension semantics attributes must
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s
• Acc ref must be at least 2D, got shape {shape}

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