jax-ml/jax · error · NotImplementedError

Non-trivial windowing is not supported for grid-free pallas_

Error message

Non-trivial windowing is not supported for grid-free pallas_call.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:1182

        arg_types=arg_types,
        num_grid=num_grid,
        get_jaxpr_indices=get_jaxpr_indices,
        ctx_factory=ctx_factory,
        dynamic_shape_replacement_enabled=dynamic_shape_replacement_enabled,
    )
  func_op.attributes["tpu.core_type"] = ir.Attribute.parse(
      f"#tpu.core_type<{kernel_type}>"
  )
  module.body.append(func_op)
  assert name not in sym_tab, f"Function name {name} already exists in symbol table."
  sym_tab.insert(func_op)
  window_params = []
  static_grid = None
  grid = mosaic_grid_mapping.grid
  if not grid and any(
      not bm.has_trivial_window() for bm in grid_mapping.block_mappings
  ):
    raise NotImplementedError(
        "Non-trivial windowing is not supported for grid-free pallas_call."
    )
  if grid:
    for i, bm in enumerate(grid_mapping.block_mappings):
      func_name = f"transform_{i}"
      # ANY and SEMAPHORE operands don't support windowing and require empty window_params.
      block_memory_space = bm.block_aval.memory_space
      if block_memory_space is None:
        block_memory_space = pallas_core.MemorySpace.DEFAULT
      tpu_memory_space = tpu_core.memory_space_to_tpu_memory_space(
          block_memory_space, kernel_type
      )
      if (
          tpu_memory_space is ANY
          or tpu_memory_space == tpu_core.MemorySpace.HBM
          or tpu_memory_space == tpu_core.MemorySpace.SEMAPHORE
      ):
        # We checked above that the block does not require windowing.

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
pallas_call(f, outspec, grid=None)  # windowed BlockSpecs present

# after
pallas_call(f, outspec, grid=(1,))  # or tiles: grid=(math.ceil(n/tile),)

Defensive patterns

Strategy: validation

Validate before calling

def safe_to_run_gridless(grid_mapping):
    return not (not grid_mapping.grid and any(
        not bm.has_trivial_window() for bm in grid_mapping.block_mappings))

Type guard

def is_gridless_compatible(gm) -> bool:
    return gm.grid or all(bm.has_trivial_window() for bm in gm.block_mappings)

Try / catch

try:
    pallas_call(f, out_spec, grid=None)
except NotImplementedError as e:
    if 'grid-free' in str(e):
        result = pallas_call(f, out_spec, grid=(1,))(...)  # retry with trivial grid

Prevention

Always pass an explicit grid, even (1,), when using windowed BlockSpecs
Keep grid-free kernels limited to full-array trivial blocks

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f4ef78908c635b94. Report an issue: GitHub.