jax-ml/jax · error · NotImplementedError

dynamic grid bounds not supported in the Mosaic GPU backend

Error message

dynamic grid bounds not supported in the Mosaic GPU backend

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/pallas_call_registration.py:58

    ctx: mlir.LoweringRuleContext,
    *args,
    jaxpr: jax_core.Jaxpr,
    interpret: bool,
    debug: bool,
    input_output_aliases: tuple[tuple[int, int], ...],
    grid_mapping: pallas_core.GridMapping,
    mesh: pallas_core.Mesh | None,
    compiler_params: pallas_core.CompilerParams | None,
    cost_estimate: pallas_core.CostEstimate | None,
    out_avals: tuple[jax_core.AbstractValue, ...],
    metadata: frozen_dict.FrozenDict[str, str] | None,
    name: str | None,
):
  del metadata, name  # TODO(sharadmv): Add metadata to HLO.
  debug_info = jaxpr.debug_info
  del interpret, out_avals
  if grid_mapping.num_dynamic_grid_bounds:
    raise NotImplementedError(
        "dynamic grid bounds not supported in the Mosaic GPU backend"
    )

  if mesh is not None and not isinstance(mesh, gpu_core.Mesh):
    raise NotImplementedError(
        f"Mesh {mesh} is not supported by the Mosaic GPU backend"
    )

  if debug:
    print(f"\nThe kernel jaxpr for pallas_call {debug_info.func_src_info}:")
    print(jaxpr)
    print(f"The grid mapping for pallas_call {debug_info.func_src_info}:")
    print(grid_mapping)

  mgpu.dialect.register_dialect(ctx.module_context.context)

  if compiler_params is None:
    gpu_params = gpu_core.CompilerParams()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all grid dimensions concrete Python ints before calling the kernel
2. Hoist grid computation outside jax-traced code so values are static
3. Pad/round shapes so a static grid can be used with in-bounds masking

Example fix

# before
grid = (num_blocks,)  # tracer
kernel[grid](...)
# after
assert isinstance(num_blocks, int)
kernel[(int(num_blocks),)](...)

assert all(isinstance(g, int) and g > 0 for g in grid), f"grid must be static ints: {grid}"

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

When it happens

Trigger: 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.

Common situations: Porting TPU pallas kernels that used dynamic grids; computing grid from symbolic batch dimensions inside jit without specializing.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Sparse metadata format not implemented for {operand_dtype=}

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