jax-ml/jax · error · NotImplementedError

Dynamic grid bounds not (yet) supported in GPU interpret mod

Error message

Dynamic grid bounds not (yet) supported in GPU interpret mode.

What it means

The GPU interpreter only handles static (Python-int) grid dimensions. If the GridMapping declares dynamic grid bounds (grid sizes computed from traced values at call time), interpret mode refuses to run.

Source

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

from jax._src.pallas.mosaic_gpu.interpret import jaxpr_interpret
from jax._src.pallas.mosaic_gpu.interpret import shared_memory as memory
from jax._src.pallas.mosaic_gpu.interpret.params import InterpretGPUParams
from jax._src.state import types as state_types
from jax._src.typing import Array
from jax._src.util import (safe_zip, split_list)


def get_races() -> gpu_callbacks.RaceDetectionState:
  return gpu_callbacks.get_races()


def reset_gpu_interpret_mode_state():
  gpu_callbacks.reset_gpu_interpret_mode_state()


def _get_grid_bounds(grid_mapping: pallas_core.GridMapping) -> tuple[int, ...]:
  if grid_mapping.num_dynamic_grid_bounds > 0:
    raise NotImplementedError(
        "Dynamic grid bounds not (yet) supported in GPU interpret mode."
    )
  result = []
  for x in grid_mapping.grid:
    # We have already tested for the absence of dynamic grid bounds. So all
    # entries in the grid should be ints.
    assert isinstance(x, int)
    result.append(x)
  return tuple(result)


def _get_grid_and_cluster_dims_and_num_threads(
    grid_mapping: pallas_core.GridMapping, mesh: mosaic_gpu_core.Mesh | None
) -> tuple[tuple[int, ...], tuple[int, ...], int]:
  if not mesh:
    num_threads = 1
    cluster_dims = ()
    grid_dims = _get_grid_bounds(grid_mapping)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a static tuple grid (plain ints) instead of a callable/dynamic grid when running in interpret mode
2. Materialize shapes before constructing the kernel invocation
3. Run on device where dynamic grids are supported
4. Update jax in case dynamic-grid interpret support landed

Example fix

# before
kernel = pallas_call(fn, out_shape, grid=lambda x: (x.shape[0],))
# after
grid = (x.shape[0],)  # static python ints
kernel = pallas_call(fn, out_shape, grid=grid)

from jax._src.pallas import pallas_core
assert grid_mapping.num_dynamic_grid_bounds == 0, 'use a static grid for interpret mode'

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

When it happens

Trigger: Creating a pallas kernel whose grid is a function of traced arguments (dynamic grid), e.g. grid=lambda a: (a.shape[0]//bs,) with traced shapes, then running in GPU interpret mode.

Common situations: Kernels using call-time grid lambdas over dynamic shapes; jit-compiled wrappers where shapes become traced; moving TPU kernels with dynamic grids to GPU interpret tests.

Related errors

• Dynamic grid bounds not (yet) supported on GPU
• 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/d6867a1e653498aa. Report an issue: GitHub.