jax-ml/jax · error · ValueError

The iteration bound corresponding to the core-parallel dimen

Error message

The iteration bound corresponding to the core-parallel dimension be the same across all subkernels.

What it means

Different subkernels in the same Mosaic kernel declare core_parallel iteration bounds with different static sizes; all subkernels must agree.

Source

Thrown at jax/_src/tpu_custom_call.py:597

    ):
      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:
    if op.operation.name != "func.func":
      continue

    if (
        "iteration_bounds" not in op.attributes
        or "dimension_semantics" not in op.attributes
    ):
      continue

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the core-parallel bound identical across all subkernels
2. Recompile subkernels for the same core count

sizes = {get_core_parallel_dim_size(sk) for sk in subkernels}
assert len(sizes) <= 1

When it happens

Trigger: Composite kernels where one subkernel uses core dim size N and another uses M != N.

Common situations: Fusing subkernels written for different core counts.

Related errors

• A single Mosaic subkernel cannot contain multiple core shard
• 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/b26951bf0bdab704. Report an issue: GitHub.