{"record":{"id":"0a8ce1721a0b3ce9","repo":"jax-ml/jax","slug":"too-many-dynamic-shapes-in-the-input-mosaic-curre","errorCode":null,"errorMessage":"Too many dynamic shapes in the input. Mosaic currently only supports up to 128 dynamic dimension values.","messagePattern":"Too many dynamic shapes in the input\\. Mosaic currently only supports up to 128 dynamic dimension values\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":187,"sourceCode":"\n  def __init__(self):\n    self.dim_expr_to_placeholder: dict[shape_poly._DimExpr, Any] = {}\n    self.placeholder_to_dim_expr: dict[Any, shape_poly._DimExpr] = {}\n\n  def snapshot(self) -> frozenset[tuple[shape_poly._DimExpr, Any]]:\n    return frozenset(self.dim_expr_to_placeholder.items())\n\n  def to_placeholder(self, dim_expr: Any) -> ir.Value:\n    if jax_core.is_constant_dim(dim_expr):\n      # avoid ints, these are not dynamic\n      return dim_expr\n    if dim_expr not in self.dim_expr_to_placeholder:\n      next_val = DIM_UPPER_BOUND - len(self.dim_expr_to_placeholder)\n      if next_val < DIM_LOWER_BOUND:\n        # In practice, even with the largest of programs, we see rarely see\n        # anything even close to this limit. It is arbitrary, and can be safely\n        # increased if needed.\n        raise ValueError(\n            \"Too many dynamic shapes in the input. Mosaic currently only\"\n            \" supports up to 128 dynamic dimension values.\"\n        )\n      self.dim_expr_to_placeholder[dim_expr] = next_val\n      # Reverse mapping - this is consumed to generate a table that is either\n      # input<>placeholder or intermediary computation<>placeholder.\n      self.placeholder_to_dim_expr[next_val] = dim_expr\n    return self.dim_expr_to_placeholder[dim_expr]\n\n\nDynamicShapeReplacementFn = Callable[\n    [tuple[jax_core.DimSize, ...]], tuple[int, ...]\n]\n\n\n@dataclasses.dataclass\nclass LoweringContext:\n  grid_sizes: tuple[int, ...]  # Includes both user and vmap axes.","sourceCodeStart":169,"sourceCodeEnd":205,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L169-L205","documentation":"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.","triggerScenarios":"A kernel whose inputs/intermediates contain more than 128 distinct symbolic dimension expressions, exhausted in to_placeholder during dynamic-shape lowering.","commonSituations":"Heavily parameterized kernels with many unrelated dynamic dims; loops unrolling into many distinct dim expressions; generated code creating unique dim expressions per iteration.","solutions":["Reduce the number of distinct dynamic dimensions (reuse dim expressions, e.g. shared size variables, instead of fresh sym expressions)","Make some dimensions static (linspace/specialize common values)","If genuinely needed, raise DIM_LOWER_BOUND/DIM_UPPER_BOUND in a local fork (the code notes the limit is arbitrary)"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"distinct_dims = collect_dynamic_dim_exprs(jaxpr)  # custom walk\nassert len(distinct_dims) <= 128, f'{len(distinct_dims)} dynamic dims exceeds Mosaic placeholder limit (128)'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Reuse one symbolic dim expression for equal dynamic sizes instead of creating new ones","Specialize frequently-repeated dimension values as static shapes"],"tags":["jax","pallas","mosaic","dynamic-shapes","limit-exceeded","lowering"],"backgroundTag":"resource-limit-exceeded","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}