{"record":{"id":"eab751ee453d08ba","repo":"jax-ml/jax","slug":"unsupported-mesh-type-type-mesh","errorCode":null,"errorMessage":"Unsupported mesh type: {type(mesh)}","messagePattern":"Unsupported mesh type: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py","lineNumber":72,"sourceCode":"    # 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)\n  elif isinstance(mesh, mosaic_gpu_core.Mesh):\n    num_threads = int(mesh.num_threads or 1)\n    cluster_dims = tuple(mesh.cluster) if mesh.cluster is not None else ()\n    grid_dims = tuple(mesh.grid)\n  else:\n    raise ValueError(f\"Unsupported mesh type: {type(mesh)}\")\n\n  reconstructed_grid = grid_dims + cluster_dims + (num_threads,)\n  if math.prod(_get_grid_bounds(grid_mapping)) != math.prod(reconstructed_grid):\n    raise NotImplementedError(\n        f\"Invalid grid {grid_mapping.grid} in grid_mapping: expected grid to\"\n        f\" have the same size as {reconstructed_grid}\"\n    )\n\n  return grid_dims, cluster_dims, num_threads\n\n\ndef _allocate_buffers_for_inputs(\n    token: jax.Array,\n    device: memory.Device,\n    invars: Sequence[Any],\n    inputs: Sequence[jax.Array],\n) -> tuple[jax.Array, list[jax.Array]]:\n  \"\"\"Allocates `GMEM` buffers for the `inputs` of a `pallas_call`.\"\"\"","sourceCodeStart":54,"sourceCodeEnd":90,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/interpret/interpret_pallas_call.py#L54-L90","documentation":"When interpreting a Mosaic GPU kernel, the interpreter derives grid/cluster/thread dimensions from the kernel's mesh. It only understands mosaic_gpu_core.Mesh (and the no-mesh case); any other mesh type is rejected.","triggerScenarios":"Passing a grid mapping whose mesh is of an unexpected type (not None, not a GridMapping-derived int grid, not mosaic_gpu_core.Mesh) — usually from API misuse or version mismatch.","commonSituations":"Mixing TPU pallas Mesh types with Mosaic GPU kernels; stale jax versions where mesh classes moved; custom grid wrapper objects.","solutions":["Pass the mesh as mosaic_gpu_core.Mesh (or None with a plain int grid)","Ensure jax and any mosaic plugin versions match","Check the printed type in the message to find which object leaked through"],"exampleFix":"# before\ngrid = some_tpu_mesh  # wrong mesh type\n# after\nfrom jax._src.pallas.mosaic_gpu.core import Mesh\ngrid = Mesh(grid=(g,), thread_name='tid', num_threads=128)","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def is_supported_mesh(mesh) -> bool:\n    from jax._src.pallas.mosaic_gpu import core\n    return mesh is None or isinstance(mesh, core.Mesh)","tryCatchPattern":null,"preventionTips":["Always construct meshes with mosaic_gpu_core.Mesh for GPU kernels"],"tags":["pallas","mosaic-gpu","mesh","interpret-mode","type-validation"],"backgroundTag":"unsupported-operation-in-interpreter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}