jax-ml/jax · error · ValueError

A single Mosaic subkernel cannot contain multiple core shard

Error message

A single Mosaic subkernel cannot contain multiple core sharding dimensions.

What it means

A subkernel declares more than one core_parallel dimension; only a single core-sharding dimension per subkernel is allowed.

Source

Thrown at jax/_src/tpu_custom_call.py:589

          "The iteration bounds and dimension semantics attributes must have"
          " the same number of elements."
      )

    subkernel_core_dim_size = None

    for dim_idx, (dim_size, dim_sem) in enumerate(
        zip(iter_bounds, dim_semantics)
    ):
      if str(dim_sem) != "#tpu.dimension_semantics<core_parallel>":
        continue

      if ir.ShapedType.is_dynamic_size(dim_size):
        raise ValueError(
            "The iteration bound corresponding to the core-parallel dimension "
            f"{dim_idx} must be statically known."
        )
      if subkernel_core_dim_size is not None:
        raise ValueError(
            "A single Mosaic subkernel cannot contain multiple core sharding "
            "dimensions."
        )
      if (
          other_subkernel_core_dim_size is not None
          and other_subkernel_core_dim_size != dim_size
      ):
        raise ValueError(
            "The iteration bound corresponding to the core-parallel dimension "
            "be the same across all subkernels."
        )
      subkernel_core_dim_size = dim_size

    return subkernel_core_dim_size

  core_parallel_dim_size = None

  for op in module.body.operations:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep only one core_parallel dimension; move other dims to arith/layout semantics

assert sum(str(d) == '#tpu.dimension_semantics<core_parallel>' for d in dim_semantics) <= 1

When it happens

Trigger: Annotating two dimensions as #tpu.dimension_semantics<core_parallel> in the same subkernel.

Common situations: Hand-written or migrated Mosaic kernels that shard across two core axes.

Related errors

• The iteration bounds and dimension semantics attributes must
• The iteration bound corresponding to the core-parallel dimen
• Invalid memory space: {memory_space!r}
• {num_queries_per_block=} must be positive.
• {vmem_limit_bytes=} must be positive.

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