{"record":{"id":"1307769a3d79aa34","repo":"jax-ml/jax","slug":"dimensions-with-parallel-semantics-must-form-a-pre","errorCode":null,"errorMessage":"Dimensions with parallel semantics must form a prefix of the grid.","messagePattern":"Dimensions with parallel semantics must form a prefix of the grid\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":1822,"sourceCode":"  Args:\n    mosaic_params: The compiler params for the Mosaic TPU backend.\n    num_dimensions_in_grid: The number of dimensions in the grid.\n\n  Returns:\n    A tuple of booleans where the entry at index `i` is `True` precisely if the\n    `i`-th dimension in the grid has parallel semantics.\n\n  Raises:\n    ValueError: If the dimensions with parallel semantics do not form a prefix\n      of the grid.\n  \"\"\"\n  if mosaic_params.dimension_semantics is None:\n    return (False,) * num_dimensions_in_grid\n  result = tuple(ds in ('parallel', mosaic_core.PARALLEL)\n                 for ds in mosaic_params.dimension_semantics)\n  for ds0, ds1 in zip(result[:-1], result[1:]):\n    if ds1 and not ds0:\n      raise ValueError(\n          'Dimensions with parallel semantics must form a prefix of the grid.'\n      )\n  return result\n\n\ndef _get_parallel_subgrid_size(\n    parallel_semantics_per_dim: tuple[bool, ...], grid: tuple[int, ...]\n) -> int:\n  \"\"\"Returns the size of the subgrid along the parallel dimensions.\"\"\"\n  return math.prod(\n      dim_size if parallel_dim else 1\n      for dim_size, parallel_dim in zip(grid, parallel_semantics_per_dim)\n  )\n\n_GridPointCoordinatesPerDim = tuple[Array, ...]\n\ndef _get_randomized_grid_coordinates(\n    grid: tuple[int, ...],","sourceCodeStart":1804,"sourceCodeEnd":1840,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L1804-L1840","documentation":"For randomized TPU interpret mode (multi-core simulation), the grid's dimension_semantics must have all 'parallel' dimensions as a prefix — no 'parallel' dimension may follow a non-parallel one. Otherwise the interpreter cannot compute consistent randomized grid coordinates and raises this ValueError.","triggerScenarios":"Calling interpret_pallas_call with CompilerParams(dimension_semantics=['parallel','loop','parallel']) or any layout where a parallel dim comes after a non-parallel dim.","commonSituations":"Matmul-style kernels with contraction dims marked loop followed by parallel broadcast dims; reordering grid dims for performance while keeping old semantics labels; multi-core interpret runs of collective kernels.","solutions":["Reorder the grid so parallel dimensions come first and relabel dimension_semantics accordingly","Mark the trailing dimensions as 'loop' if they carry sequential dependencies","Permute grid axes (and corresponding BlockSpec index_maps) to satisfy the prefix rule"],"exampleFix":"# before\ndimension_semantics=['parallel', 'loop', 'parallel']\n# after\ndimension_semantics=['parallel', 'parallel', 'loop']  # reorder grid accordingly","handlingStrategy":"validation","validationCode":"def semantics_ok(ds):\n    seen_non_parallel = False\n    for d in ds:\n        if d != 'parallel':\n            seen_non_parallel = True\n        elif seen_non_parallel:\n            return False\n    return True\nassert semantics_ok(dimension_semantics)","typeGuard":"def is_parallel_prefix(ds) -> bool:\n    p = [d in ('parallel',) for d in ds]\n    return all(p[:sum(p)]) and not any(p[sum(p):])","tryCatchPattern":"try:\n    interpret_run(kernel)\nexcept ValueError as e:\n    if 'prefix of the grid' in str(e):\n        # reorder grid dims so parallel dims come first, update index_maps\n        raise","preventionTips":["Design grids with parallel dims first","Update dimension_semantics whenever reordering the grid","Validate semantics labels in unit tests for multi-core interpret runs"],"tags":["jax","pallas","tpu","grid","dimension-semantics"],"backgroundTag":"invalid-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}