{"record":{"id":"6f121d7e09807ae4","repo":"jax-ml/jax","slug":"invalid-grid-grid-mapping-grid-in-grid-mapping","errorCode":null,"errorMessage":"Invalid grid {grid_mapping.grid} in grid_mapping: expected grid to have the same size as {reconstructed_grid}","messagePattern":"Invalid grid (.+?) in grid_mapping: expected grid to have the same size as (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py","lineNumber":76,"sourceCode":"\n\ndef _get_grid_and_cluster_dims_and_num_threads(\n    grid_mapping: pallas_core.GridMapping, mesh: mosaic_gpu_core.Mesh | None\n) -> tuple[tuple[int, ...], tuple[int, ...], int]:\n  if not mesh:\n    num_threads = 1\n    cluster_dims = ()\n    grid_dims = _get_grid_bounds(grid_mapping)\n  elif isinstance(mesh, mosaic_gpu_core.Mesh):\n    num_threads = int(mesh.num_threads or 1)\n    cluster_dims = tuple(mesh.cluster) if mesh.cluster is not None else ()\n    grid_dims = tuple(mesh.grid)\n  else:\n    raise ValueError(f\"Unsupported mesh type: {type(mesh)}\")\n\n  reconstructed_grid = grid_dims + cluster_dims + (num_threads,)\n  if math.prod(_get_grid_bounds(grid_mapping)) != math.prod(reconstructed_grid):\n    raise NotImplementedError(\n        f\"Invalid grid {grid_mapping.grid} in grid_mapping: expected grid to\"\n        f\" have the same size as {reconstructed_grid}\"\n    )\n\n  return grid_dims, cluster_dims, num_threads\n\n\ndef _allocate_buffers_for_inputs(\n    token: jax.Array,\n    device: memory.Device,\n    invars: Sequence[Any],\n    inputs: Sequence[jax.Array],\n) -> tuple[jax.Array, list[jax.Array]]:\n  \"\"\"Allocates `GMEM` buffers for the `inputs` of a `pallas_call`.\"\"\"\n  # TODO(nrink): This code is a simplified version to the corresponding TPU\n  # interpreter code. Eventually, we should merge the two.\n  input_buffer_keys = []\n  for var, value in safe_zip(invars, inputs):","sourceCodeStart":58,"sourceCodeEnd":94,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py#L58-L94","documentation":"Consistency check: the product of the grid mapping's dimensions must equal the product of the reconstructed (grid × cluster × num_threads) dims from the mesh. A mismatch means the mesh and the declared grid describe different amounts of parallelism.","triggerScenarios":"Declaring a pallas_call grid whose total size differs from grid*cluster*threads of the Mosaic Mesh — e.g. grid=(8,) but Mesh(grid=(4,), num_threads=128) (or similar product mismatch).","commonSituations":"Editing one of grid/mesh/num_threads without updating the others; copy-paste kernels with mismatched launch configs; migrating kernels between thread counts.","solutions":["Make the products match: ensure prod(grid_mapping.grid) == prod(mesh.grid) * prod(cluster) * num_threads","Usually set the pallas_call grid from the mesh: grid=tuple(mesh.grid) (times cluster dims)","Recompute launch dims after changing num_threads","Check the printed expected tuple in the message and align to it"],"exampleFix":"# before\nkernel = pallas_call(fn, out, grid=(8,))\nkernel(..., grid=Mesh(grid=(4,), num_threads=128, ...))  # mismatch\n# after\nmesh = Mesh(grid=(4,), num_threads=128, thread_name='tid')\nkernel = pallas_call(fn, out, grid=tuple(mesh.grid))","handlingStrategy":"validation","validationCode":"import math\nexpected = tuple(mesh.grid) + (tuple(mesh.cluster) if mesh.cluster else ()) + (int(mesh.num_threads or 1),)\nassert math.prod(grid_mapping.grid) == math.prod(expected)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive pallas_call grid from the Mesh to keep products consistent","Re-validate launch dims after changing num_threads"],"tags":["pallas","mosaic-gpu","grid","mesh","config-mismatch"],"backgroundTag":"grid-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}