{"record":{"id":"d6867a1e653498aa","repo":"jax-ml/jax","slug":"dynamic-grid-bounds-not-yet-supported-in-gpu-int","errorCode":null,"errorMessage":"Dynamic grid bounds not (yet) supported in GPU interpret mode.","messagePattern":"Dynamic grid bounds not \\(yet\\) supported in GPU interpret mode\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py","lineNumber":48,"sourceCode":"from jax._src.pallas.mosaic_gpu.interpret import jaxpr_interpret\nfrom jax._src.pallas.mosaic_gpu.interpret import shared_memory as memory\nfrom jax._src.pallas.mosaic_gpu.interpret.params import InterpretGPUParams\nfrom jax._src.state import types as state_types\nfrom jax._src.typing import Array\nfrom jax._src.util import (safe_zip, split_list)\n\n\ndef get_races() -> gpu_callbacks.RaceDetectionState:\n  return gpu_callbacks.get_races()\n\n\ndef reset_gpu_interpret_mode_state():\n  gpu_callbacks.reset_gpu_interpret_mode_state()\n\n\ndef _get_grid_bounds(grid_mapping: pallas_core.GridMapping) -> tuple[int, ...]:\n  if grid_mapping.num_dynamic_grid_bounds > 0:\n    raise NotImplementedError(\n        \"Dynamic grid bounds not (yet) supported in GPU interpret mode.\"\n    )\n  result = []\n  for x in grid_mapping.grid:\n    # We have already tested for the absence of dynamic grid bounds. So all\n    # entries in the grid should be ints.\n    assert isinstance(x, int)\n    result.append(x)\n  return tuple(result)\n\n\ndef _get_grid_and_cluster_dims_and_num_threads(\n    grid_mapping: pallas_core.GridMapping, mesh: mosaic_gpu_core.Mesh | None\n) -> tuple[tuple[int, ...], tuple[int, ...], int]:\n  if not mesh:\n    num_threads = 1\n    cluster_dims = ()\n    grid_dims = _get_grid_bounds(grid_mapping)","sourceCodeStart":30,"sourceCodeEnd":66,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py#L30-L66","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pass a static tuple grid (plain ints) instead of a callable/dynamic grid when running in interpret mode","Materialize shapes before constructing the kernel invocation","Run on device where dynamic grids are supported","Update jax in case dynamic-grid interpret support landed"],"exampleFix":"# before\nkernel = pallas_call(fn, out_shape, grid=lambda x: (x.shape[0],))\n# after\ngrid = (x.shape[0],)  # static python ints\nkernel = pallas_call(fn, out_shape, grid=grid)","handlingStrategy":"validation","validationCode":"from jax._src.pallas import pallas_core\nassert grid_mapping.num_dynamic_grid_bounds == 0, 'use a static grid for interpret mode'","typeGuard":"def is_static_grid(grid) -> bool:\n    return isinstance(grid, tuple) and all(isinstance(g, int) for g in grid)","tryCatchPattern":null,"preventionTips":["Pass int-tuple grids for interpret runs","Materialize shapes before kernel creation"],"tags":["pallas","mosaic-gpu","interpret-mode","dynamic-grid","not-implemented"],"backgroundTag":"traced-value-where-static-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}