{"record":{"id":"e6d84f0cc9e2a11c","repo":"jax-ml/jax","slug":"dynamic-grid-bounds-not-supported-in-the-mosaic-gp-e6d84f","errorCode":null,"errorMessage":"dynamic grid bounds not supported in the Mosaic GPU backend","messagePattern":"dynamic grid bounds not supported in the Mosaic GPU backend","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/pallas_call_registration.py","lineNumber":58,"sourceCode":"    ctx: mlir.LoweringRuleContext,\n    *args,\n    jaxpr: jax_core.Jaxpr,\n    interpret: bool,\n    debug: bool,\n    input_output_aliases: tuple[tuple[int, int], ...],\n    grid_mapping: pallas_core.GridMapping,\n    mesh: pallas_core.Mesh | None,\n    compiler_params: pallas_core.CompilerParams | None,\n    cost_estimate: pallas_core.CostEstimate | None,\n    out_avals: tuple[jax_core.AbstractValue, ...],\n    metadata: frozen_dict.FrozenDict[str, str] | None,\n    name: str | None,\n):\n  del metadata, name  # TODO(sharadmv): Add metadata to HLO.\n  debug_info = jaxpr.debug_info\n  del interpret, out_avals\n  if grid_mapping.num_dynamic_grid_bounds:\n    raise NotImplementedError(\n        \"dynamic grid bounds not supported in the Mosaic GPU backend\"\n    )\n\n  if mesh is not None and not isinstance(mesh, gpu_core.Mesh):\n    raise NotImplementedError(\n        f\"Mesh {mesh} is not supported by the Mosaic GPU backend\"\n    )\n\n  if debug:\n    print(f\"\\nThe kernel jaxpr for pallas_call {debug_info.func_src_info}:\")\n    print(jaxpr)\n    print(f\"The grid mapping for pallas_call {debug_info.func_src_info}:\")\n    print(grid_mapping)\n\n  mgpu.dialect.register_dialect(ctx.module_context.context)\n\n  if compiler_params is None:\n    gpu_params = gpu_core.CompilerParams()","sourceCodeStart":40,"sourceCodeEnd":76,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/pallas_call_registration.py#L40-L76","documentation":"pallas_call lowering for the Mosaic GPU backend does not support dynamic grid bounds; the grid must be fully static (known shapes/constants) at trace time.","triggerScenarios":"Passing a grid containing dynamic extents (e.g. computed from abstract values with jax.make_jaxpr, or grid entries that are tracers/unknown) to a pallas_call targeting Mosaic GPU.","commonSituations":"Porting TPU pallas kernels that used dynamic grids; computing grid from symbolic batch dimensions inside jit without specializing.","solutions":["Make all grid dimensions concrete Python ints before calling the kernel","Hoist grid computation outside jax-traced code so values are static","Pad/round shapes so a static grid can be used with in-bounds masking"],"exampleFix":"# before\ngrid = (num_blocks,)  # tracer\nkernel[grid](...)\n# after\nassert isinstance(num_blocks, int)\nkernel[(int(num_blocks),)](...)","handlingStrategy":"validation","validationCode":"assert all(isinstance(g, int) and g > 0 for g in grid), f\"grid must be static ints: {grid}\"","typeGuard":"def is_static_grid(grid) -> bool:\n    return all(isinstance(g, int) for g in grid)","tryCatchPattern":null,"preventionTips":["Compute grids in untraced Python code","Round shapes up to static sizes with masking"],"tags":["pallas","mosaic-gpu","grid","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}