{"record":{"id":"f4ef78908c635b94","repo":"jax-ml/jax","slug":"non-trivial-windowing-is-not-supported-for-grid-fr","errorCode":null,"errorMessage":"Non-trivial windowing is not supported for grid-free pallas_call.","messagePattern":"Non-trivial windowing is not supported for grid-free pallas_call\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":1182,"sourceCode":"        arg_types=arg_types,\n        num_grid=num_grid,\n        get_jaxpr_indices=get_jaxpr_indices,\n        ctx_factory=ctx_factory,\n        dynamic_shape_replacement_enabled=dynamic_shape_replacement_enabled,\n    )\n  func_op.attributes[\"tpu.core_type\"] = ir.Attribute.parse(\n      f\"#tpu.core_type<{kernel_type}>\"\n  )\n  module.body.append(func_op)\n  assert name not in sym_tab, f\"Function name {name} already exists in symbol table.\"\n  sym_tab.insert(func_op)\n  window_params = []\n  static_grid = None\n  grid = mosaic_grid_mapping.grid\n  if not grid and any(\n      not bm.has_trivial_window() for bm in grid_mapping.block_mappings\n  ):\n    raise NotImplementedError(\n        \"Non-trivial windowing is not supported for grid-free pallas_call.\"\n    )\n  if grid:\n    for i, bm in enumerate(grid_mapping.block_mappings):\n      func_name = f\"transform_{i}\"\n      # ANY and SEMAPHORE operands don't support windowing and require empty window_params.\n      block_memory_space = bm.block_aval.memory_space\n      if block_memory_space is None:\n        block_memory_space = pallas_core.MemorySpace.DEFAULT\n      tpu_memory_space = tpu_core.memory_space_to_tpu_memory_space(\n          block_memory_space, kernel_type\n      )\n      if (\n          tpu_memory_space is ANY\n          or tpu_memory_space == tpu_core.MemorySpace.HBM\n          or tpu_memory_space == tpu_core.MemorySpace.SEMAPHORE\n      ):\n        # We checked above that the block does not require windowing.","sourceCodeStart":1164,"sourceCodeEnd":1200,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L1164-L1200","documentation":"Pallas supports 'windowing' (non-trivial block windows with padding) only in combination with an explicit grid. If a pallas_call has no grid (grid-free) but any BlockMapping has a non-trivial window, the lowering raises NotImplementedError because there is no iteration structure over which windows could be evaluated.","triggerScenarios":"Calling pallas_call with grid=None (or an empty grid) while a BlockSpec/index_map implies a non-trivial window, e.g. block shapes smaller than the array with offsets or padding, or using experimental windowed BlockSpecs without a grid.","commonSituations":"Converting a gridded kernel to a single-shot grid-free kernel but keeping windowed BlockSpecs; using autodiff or transpose that generates windowed block mappings for a grid-free call; copy-pasting windowed examples without the accompanying grid argument.","solutions":["Pass an explicit grid (e.g. grid=(1,) or the tile counts) to the pallas_call","Make all block mappings use trivial windows (full-shape blocks with all-zero index maps) if you truly want grid-free execution","Rewrite the kernel to slice inputs manually inside the body instead of relying on windowing"],"exampleFix":"# before\npallas_call(f, outspec, grid=None)  # windowed BlockSpecs present\n\n# after\npallas_call(f, outspec, grid=(1,))  # or tiles: grid=(math.ceil(n/tile),)","handlingStrategy":"validation","validationCode":"def safe_to_run_gridless(grid_mapping):\n    return not (not grid_mapping.grid and any(\n        not bm.has_trivial_window() for bm in grid_mapping.block_mappings))","typeGuard":"def is_gridless_compatible(gm) -> bool:\n    return gm.grid or all(bm.has_trivial_window() for bm in gm.block_mappings)","tryCatchPattern":"try:\n    pallas_call(f, out_spec, grid=None)\nexcept NotImplementedError as e:\n    if 'grid-free' in str(e):\n        result = pallas_call(f, out_spec, grid=(1,))(...)  # retry with trivial grid","preventionTips":["Always pass an explicit grid, even (1,), when using windowed BlockSpecs","Keep grid-free kernels limited to full-array trivial blocks"],"tags":["jax","pallas","tpu","grid","windowing"],"backgroundTag":"pallas-unsupported-windowing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}