{"record":{"id":"402ff961dd3e4612","repo":"jax-ml/jax","slug":"the-iteration-bound-corresponding-to-the-core-para","errorCode":null,"errorMessage":"The iteration bound corresponding to the core-parallel dimension {dim_idx} must be statically known.","messagePattern":"The iteration bound corresponding to the core-parallel dimension (.+?) must be statically known\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/tpu_custom_call.py","lineNumber":584,"sourceCode":"      iter_bounds: ir.DenseI64ArrayAttr,\n      other_subkernel_core_dim_size: int | None = None) -> int | None:\n\n    if len(iter_bounds) != len(dim_semantics):\n      raise ValueError(\n          \"The iteration bounds and dimension semantics attributes must have\"\n          \" the same number of elements.\"\n      )\n\n    subkernel_core_dim_size = None\n\n    for dim_idx, (dim_size, dim_sem) in enumerate(\n        zip(iter_bounds, dim_semantics)\n    ):\n      if str(dim_sem) != \"#tpu.dimension_semantics<core_parallel>\":\n        continue\n\n      if ir.ShapedType.is_dynamic_size(dim_size):\n        raise ValueError(\n            \"The iteration bound corresponding to the core-parallel dimension \"\n            f\"{dim_idx} must be statically known.\"\n        )\n      if subkernel_core_dim_size is not None:\n        raise ValueError(\n            \"A single Mosaic subkernel cannot contain multiple core sharding \"\n            \"dimensions.\"\n        )\n      if (\n          other_subkernel_core_dim_size is not None\n          and other_subkernel_core_dim_size != dim_size\n      ):\n        raise ValueError(\n            \"The iteration bound corresponding to the core-parallel dimension \"\n            \"be the same across all subkernels.\"\n        )\n      subkernel_core_dim_size = dim_size\n","sourceCodeStart":566,"sourceCodeEnd":602,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/tpu_custom_call.py#L566-L602","documentation":"The core-parallel dimension of a subkernel has a dynamic (unknown at compile time) iteration bound; it must be static to determine core parallelism.","triggerScenarios":"A subkernel whose core_parallel dimension has an iteration bound of ShapedType dynamic size, e.g. bounds derived from a runtime-dependent value.","commonSituations":"Making loop bounds data-dependent (from a tracer or an arg with dynamic shape) in a Pallas kernel.","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"],"exampleFix":"// before\nbound = x.shape[0] if dynamic else None\n// after\nBOUND = 128  # static constant used in kernel grid","handlingStrategy":"validation","validationCode":"assert all(not is_dynamic(b) for b in core_parallel_bounds)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use static Python ints for grid/bound sizes in Pallas kernels","Avoid deriving loop bounds from runtime data"],"tags":["tpu","pallas","dynamic-shape","static-bounds","jax"],"backgroundTag":"dynamic-shape-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}