{"record":{"id":"5661ebfdac9203b5","repo":"jax-ml/jax","slug":"dynamic-indexing-not-supported-in-gpu-interpret-mo","errorCode":null,"errorMessage":"Dynamic indexing not supported in GPU interpret mode","messagePattern":"Dynamic indexing not supported in GPU interpret mode","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py","lineNumber":513,"sourceCode":"  else:\n    uninit_array[tuple(slice(s) for s in ret.shape)] = ret\n    return uninit_array\n\n\ndef _is_dynamic(indexer: indexing.NDIndexer) -> bool:\n  return any(\n      isinstance(idx, indexing.Slice)\n      and (idx.is_dynamic_start or idx.is_dynamic_size)\n      for idx in indexer.indices\n  )\n\n\ndef _validate_transforms(transforms):\n  for transform in transforms:\n    match transform:\n      case indexing.NDIndexer():\n        if _is_dynamic(transform):\n          raise ValueError(\n              \"Dynamic indexing not supported in GPU interpret mode\"\n          )\n      case _:\n        raise ValueError(f\"Unsupported transform: {transform}\")\n\n\ndef _get(\n    token: jax.Array,\n    mesh_location: memory.MeshLocation,\n    thread: memory.Thread | None,\n    allocation_key_as_array: jax.Array,\n    transforms,\n    block_indices=None,\n    grid_loop_idx=None,\n    clock=None,\n    increment_clock: bool = True,\n    source_info=None,\n    input_name=None,","sourceCodeStart":495,"sourceCodeEnd":531,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py#L495-L531","documentation":"The GPU interpret mode only supports static (compile-time known) index expressions. If an NDIndexer transform contains dynamic (traced) values the interpreter cannot compute the read range on host and raises this ValueError.","triggerScenarios":"Constructing a Block whose dimension indices are functions of runtime values (e.g. loop-carried jnp scalars) and running the kernel in GPU interpret mode; any transform where _is_dynamic() detects non-static indices.","commonSituations":"Porting a TPU Pallas kernel that used dynamic indices to Mosaic GPU; indices derived from computed scalars rather than grid mappings; switching a kernel to interpret mode when indices depend on block arguments marked dynamic.","solutions":["Make all Block indices static: derive them from grid indices / python ints instead of traced values","Use static_unsafe or hoist the dynamic value to a Python scalar before constructing the Block","Test on device (non-interpret) if dynamic indexing is genuinely required, since interpret mode cannot emulate it","Refactor to express the varying index as a grid dimension instead"],"exampleFix":"# before\nstart = some_traced_scalar  # dynamic\nblk = x[ds[start, bs], :]\n# after\nstart = int(start_value)  # python int, static per interpretation step\nblk = x[ds[start, bs], :]","handlingStrategy":"validation","validationCode":"from jax._src.pallas.mosaic_gpu.interpret.gpu_callbacks import _is_dynamic\nassert not any(_is_dynamic(t) for t in blk.transforms), 'dynamic index'","typeGuard":"def is_static_index(i) -> bool:\n    return isinstance(i, (int,)) or not isinstance(i, jax.Array) or getattr(i, 'aval', None) is None","tryCatchPattern":null,"preventionTips":["Keep Block indices as Python ints","Derive indices from grid dims, not traced values"],"tags":["pallas","mosaic-gpu","interpret-mode","dynamic-indexing"],"backgroundTag":"traced-value-where-static-required","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}