jax-ml/jax · error · ValueError

The iteration bounds and dimension semantics attributes must

Error message

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

What it means

A Mosaic subkernel's iteration bounds array and dimension semantics array have different lengths; they must be parallel 1:1.

Source

Thrown at jax/_src/tpu_custom_call.py:570

  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(
      dim_semantics: ir.ArrayAttr,
      iter_bounds: ir.DenseI64ArrayAttr,
      other_subkernel_core_dim_size: int | None = None) -> int | None:

    if len(iter_bounds) != len(dim_semantics):
      raise ValueError(
          "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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Update JAX/jaxlib to match the Mosaic version that produced the module
2. Regenerate the kernel rather than reusing stale lowered IR

assert len(in_bounds) == len(dim_semantics) for each subkernel before lowering

When it happens

Trigger: A lowered tpu.subkernel op whose in_bounds length differs from dim_semantics length — typically from malformed/incompatible Mosaic IR.

Common situations: Version mismatch between Mosaic emitter and JAX lowering; manually constructed IR.

Understand the failure class

Background: Schema validation failed / invalid input schema: payload rejected because its shape doesn't match the expected schema — this error's family across 28 libraries.

Related errors

• A single Mosaic subkernel cannot contain multiple core shard
• Invalid memory space: {memory_space!r}
• Unknown core type: {core_type}
• The iteration bound corresponding to the core-parallel dimen
• {num_queries_per_block=} must be positive.

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