jax-ml/jax · error · NotImplementedError

Dynamic grid bounds not supported in the Mosaic GPU lowering

Error message

Dynamic grid bounds not supported in the Mosaic GPU lowering.

What it means

Error "Dynamic grid bounds not supported in the Mosaic GPU lowering." thrown in jax-ml/jax.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:777

      transforms=bm.transforms,
  )


def lower_pipelined_jaxpr_to_module(
    grid_mapping: pallas_core.GridMapping,
    gpu_mesh: gpu_core.Mesh | None,
    jax_mesh: mesh_lib.Mesh | None,
    jaxpr: jax_core.Jaxpr,
    params: gpu_core.CompilerParams,
    cost_estimate: pallas_core.CostEstimate | None,
    outer_traceback: xc.Traceback | None = None,
) -> LoweringResult:
  del cost_estimate  # Unused.

  assert len(jaxpr.outvars) == 0
  assert not grid_mapping.vmapped_dims
  if grid_mapping.num_dynamic_grid_bounds:
    raise NotImplementedError(
        "Dynamic grid bounds not supported in the Mosaic GPU lowering."
    )
  if grid_mapping.num_index_operands:
    raise NotImplementedError(
        "Scalar prefetch not supported in Mosaic GPU lowering."
    )

  block_mappings = grid_mapping.block_mappings
  _check_block_mappings(block_mappings, jaxpr.debug_info)
  in_block_mappings, out_block_mappings = util.split_list(
      block_mappings, [grid_mapping.num_inputs]
  )

  grid: Sequence[int]
  if gpu_mesh:
    assert isinstance(gpu_mesh, gpu_core.Mesh)
    block = (128 * (gpu_mesh.num_threads or 1), 1, 1)
    grid = gpu_mesh.grid

View on GitHub (pinned to 1e1c6a8fc0)

When it happens

Trigger: Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:777 when the library encounters an invalid state.

Common situations: See trigger scenarios.

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