jax-ml/jax · error · NotImplementedError

Dynamic grid bounds not (yet) supported on GPU

Error message

Dynamic grid bounds not (yet) supported on GPU

What it means

interpret_pallas_call rejects kernels whose GridMapping declares dynamic grid bounds that arrive as runtime arguments (dynamic_grid_args non-empty). GPU interpret mode only supports fully static grids.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py:418

      gpu_callbacks.TOKEN_SHAPE_DTYPE,
      token,
      ordered=True,
  )

  token = gpu_callbacks.call_initialize_shared_memory(
      token=token,
      num_gpus=jnp.int32(device_info.num_devices),
      num_threads_per_block=jnp.int32(num_threads_per_block),
      num_blocks_per_cluster=jnp.int32(num_blocks_per_cluster),
      interpret_params=interpret_params,
  )

  dynamic_grid_args, scalars, inputs = split_list(
      args,
      [grid_mapping.num_dynamic_grid_bounds, grid_mapping.num_index_operands],
  )
  if dynamic_grid_args:
    raise NotImplementedError("Dynamic grid bounds not (yet) supported on GPU")
  if scalars:
    raise NotImplementedError("Scalar arguments not (yet) supported on GPU")

  assert grid_mapping.num_index_operands == 0

  token, input_buffer_keys = _allocate_buffers_for_inputs(
      token,
      device,
      jaxpr.invars[: grid_mapping.num_inputs],
      inputs,
  )

  token, output_buffers = _allocate_buffers_for_outputs(
      token,
      device,
      num_threads_per_block,
      input_output_aliases,
      grid_mapping,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a static grid for interpret-mode runs
2. Construct the grid from concrete Python ints outside jit
3. Skip interpret mode for kernels requiring dynamic grids (test on device)
4. Update jax in case support was added

Example fix

# before
kernel(x, dynamic_grid_args=(n_blocks,))  # traced n_blocks
# after
kernel(x)  # with grid=(int(n_blocks),) fixed at pallas_call creation

assert not dynamic_grid_args, 'GPU interpret mode requires a static grid'

def is_static_grid(grid) -> bool:
    return isinstance(grid, tuple) and all(isinstance(g, int) for g in grid)

When it happens

Trigger: Invoking a pallas kernel with dynamic grid bounds (grid sizes passed as traced arguments at call time) while GPU interpret mode is active.

Common situations: Debugging with interpret mode a kernel that uses call-time dynamic grids that works on device; shapes flowing through jit making grid args dynamic.

Related errors

• Dynamic grid bounds not (yet) supported in GPU interpret mod
• predicate not supported
• reduction_op not supported
• Scalar arguments not (yet) supported on GPU
• Unsupported memory space: {space}

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