{"record":{"id":"ba1029c90d6e4a47","repo":"jax-ml/jax","slug":"revisited-block-output-ranges-i-of-output-i-i","errorCode":null,"errorMessage":"Revisited block {output_ranges[i]} of output {i} in iteration {loop_idx}. The block was previously visited in iterations {past_output_ranges[past_idxs[0]][0]} through {past_output_ranges[past_idxs[-1]][0]} .","messagePattern":"Revisited block (.+?) of output (.+?) in iteration (.+?)\\. The block was previously visited in iterations (.+?) through (.+?) \\.","errorType":"validation","errorClass":"RuntimeError","httpStatus":null,"severity":"critical","filePath":"jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py","lineNumber":293,"sourceCode":"  shared_memory = _get_shared_memory()\n  past_output_ranges = shared_memory.output_ranges[(device_id, local_core_id)]\n  if not past_output_ranges:\n    past_output_ranges.append((loop_idx, output_ranges))\n    return token\n\n  for i in range(len(output_ranges)):\n    if output_ranges[i] is None:\n      continue\n    if past_output_ranges[-1][1][i] == output_ranges[i]:\n      continue\n    # TODO(jburnim): Do something constant time instead of linear here.\n    past_idxs = [\n        j\n        for j, ors in enumerate(past_output_ranges)\n        if ors[1][i] == output_ranges[i]\n    ]\n    if past_idxs:\n      raise RuntimeError(\n          f'Revisited block {output_ranges[i]} of output {i} in iteration '\n          f'{loop_idx}. The block was previously visited in iterations '\n          f'{past_output_ranges[past_idxs[0]][0]} through '\n          f'{past_output_ranges[past_idxs[-1]][0]} .'\n      )\n\n  past_output_ranges.append((loop_idx, output_ranges))\n  return token\n\n\n@fail_on_exception\ndef _validate(token, device_id):\n  device_id = int(device_id)\n\n  shared_memory = _get_shared_memory()\n  semaphores = shared_memory.get_sempahores_with_nonzero_count(device_id)\n  if semaphores:\n    sem, global_core_id = semaphores[0]","sourceCodeStart":275,"sourceCodeEnd":311,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py#L275-L311","documentation":"The TPU interpreter tracks which output blocks each grid iteration writes; writing the same output block in two different iterations is a kernel correctness bug (race/undefined ordering), so the interpreter raises RuntimeError with both the current and previously-covered iteration ranges.","triggerScenarios":"A pallas_call grid where the index_fn maps two different grid iterations to the same output block — e.g., an index map that ignores some grid index, or a block size vs grid shape mismatch causing overlapping writes.","commonSituations":"Reusing an index dimension in the index map (forgetting to multiply by block count); output shape not evenly divisible by block size so the tail block overlaps; refactored grid shapes without updating index maps.","solutions":["Fix index_fn so each grid iteration maps to a distinct output block (check multipliers like i * num_blocks)","Ensure output shape is divisible by block shape, or use masking instead of overlapping writes","Print/trace output_ranges across a tiny grid to find the colliding iterations","If overlap is intentional (e.g., reductions), restructure to accumulate in VMEM/SMEM and write once"],"exampleFix":"// before\nout_specs=BlockSpec((B,), index_fn=lambda i, j: (j,))  # collides across i\n// after\nout_specs=BlockSpec((B,), index_fn=lambda i, j: (i * num_j + j,))","handlingStrategy":"validation","validationCode":"def index_map_unique(grid, index_fn, num_outputs):\n    seen = set()\n    for idx in itertools.product(*map(range, grid)):\n        blocks = index_fn(*idx)\n        key = tuple(map(tuple, blocks))\n        if key in seen:\n            return False, idx, key\n        seen.add(key)\n    return True, None, None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Property-test index_fn injectivity on small grids in CI","Verify output shapes are divisible by block shapes","Never ignore a grid dimension in index_fn"],"tags":["jax","pallas","tpu","interpreter","overlapping-writes","grid"],"backgroundTag":"overlapping-output-writes","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}