jax-ml/jax · error · ValueError

Too many dynamic shapes in the input. Mosaic currently only

Error message

Too many dynamic shapes in the input. Mosaic currently only supports up to 128 dynamic dimension values.

What it means

Mosaic lowering ran out of placeholder slots for dynamic dimension expressions. Placeholders are drawn from a bounded integer range (DIM_LOWER_BOUND..DIM_UPPER_BOUND, 128 values); more than 128 distinct dynamic dims cannot be encoded.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:187

  def __init__(self):
    self.dim_expr_to_placeholder: dict[shape_poly._DimExpr, Any] = {}
    self.placeholder_to_dim_expr: dict[Any, shape_poly._DimExpr] = {}

  def snapshot(self) -> frozenset[tuple[shape_poly._DimExpr, Any]]:
    return frozenset(self.dim_expr_to_placeholder.items())

  def to_placeholder(self, dim_expr: Any) -> ir.Value:
    if jax_core.is_constant_dim(dim_expr):
      # avoid ints, these are not dynamic
      return dim_expr
    if dim_expr not in self.dim_expr_to_placeholder:
      next_val = DIM_UPPER_BOUND - len(self.dim_expr_to_placeholder)
      if next_val < DIM_LOWER_BOUND:
        # In practice, even with the largest of programs, we see rarely see
        # anything even close to this limit. It is arbitrary, and can be safely
        # increased if needed.
        raise ValueError(
            "Too many dynamic shapes in the input. Mosaic currently only"
            " supports up to 128 dynamic dimension values."
        )
      self.dim_expr_to_placeholder[dim_expr] = next_val
      # Reverse mapping - this is consumed to generate a table that is either
      # input<>placeholder or intermediary computation<>placeholder.
      self.placeholder_to_dim_expr[next_val] = dim_expr
    return self.dim_expr_to_placeholder[dim_expr]


DynamicShapeReplacementFn = Callable[
    [tuple[jax_core.DimSize, ...]], tuple[int, ...]
]


@dataclasses.dataclass
class LoweringContext:
  grid_sizes: tuple[int, ...]  # Includes both user and vmap axes.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reduce the number of distinct dynamic dimensions (reuse dim expressions, e.g. shared size variables, instead of fresh sym expressions)
2. Make some dimensions static (linspace/specialize common values)
3. If genuinely needed, raise DIM_LOWER_BOUND/DIM_UPPER_BOUND in a local fork (the code notes the limit is arbitrary)

distinct_dims = collect_dynamic_dim_exprs(jaxpr)  # custom walk
assert len(distinct_dims) <= 128, f'{len(distinct_dims)} dynamic dims exceeds Mosaic placeholder limit (128)'

When it happens

Trigger: A kernel whose inputs/intermediates contain more than 128 distinct symbolic dimension expressions, exhausted in to_placeholder during dynamic-shape lowering.

Common situations: Heavily parameterized kernels with many unrelated dynamic dims; loops unrolling into many distinct dim expressions; generated code creating unique dim expressions per iteration.

Related errors

• program id was requested but no grid was provided.
• Invalid axis {axis} for num_programs
• Error refining shapes. {dump_module_message(module, "before_
• Acc ref must be at least 2D, got shape {shape}
• Unsupported core type: {core_type}

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