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 {dim_idx} must be statically known.

What it means

The core-parallel dimension of a subkernel has a dynamic (unknown at compile time) iteration bound; it must be static to determine core parallelism.

Source

Thrown at jax/_src/tpu_custom_call.py:584

      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:
        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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make the core-parallel loop bound a Python int / static constant
Pad or fix the dimension size so it is known at trace time

Example fix

// before
bound = x.shape[0] if dynamic else None
// after
BOUND = 128  # static constant used in kernel grid

Defensive patterns

Strategy: validation

Validate before calling

assert all(not is_dynamic(b) for b in core_parallel_bounds)

Prevention

Use static Python ints for grid/bound sizes in Pallas kernels
Avoid deriving loop bounds from runtime data

When it happens

Trigger: A subkernel whose core_parallel dimension has an iteration bound of ShapedType dynamic size, e.g. bounds derived from a runtime-dependent value.

Common situations: Making loop bounds data-dependent (from a tracer or an arg with dynamic shape) in a Pallas kernel.

Related errors

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