{"record":{"id":"d264e70ac84c4bf3","repo":"jax-ml/jax","slug":"top-level-all-gather-works-when-all-mesh-axes-of-c","errorCode":null,"errorMessage":"top_level_all_gather works when all mesh axes of context mesh are explicit. Got {get_abstract_mesh()}","messagePattern":"top_level_all_gather works when all mesh axes of context mesh are explicit\\. Got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":2248,"sourceCode":"            f\"top_level_all_gather doesn't allow input {aval} to be unsharded\"\n            f\" on dimension {axis} when {out_spec=}.\")\n      i = i if isinstance(i, tuple) else (i,)\n      o = o if o is None or isinstance(o, tuple) else (o,)\n      if o is not None and i[:len(o)] != o:\n        raise ValueError(\n            'top_level_all_gather maintains `top_level_all_gather(x, ...) == x`'\n            f\" property. The {in_spec=} and {out_spec=} don't satisfy this\"\n            f' property. Please change your out_spec of array dimension {axis} so'\n            \" that it's a prefix of in_spec\")\n      axis_name = i if o is None else i[-len(o):]\n      x = lax_parallel.all_gather(x, axis_name=axis_name, axis=axis,\n                                  tiled=True, to='reduced')\n    return x\n  return api.jit(shard_map(f_shmap, out_specs=out_spec))(x)\n\ndef top_level_all_gather(xs, out_sharding, *, multi_dim: bool = False):\n  if not get_abstract_mesh().are_all_axes_explicit:\n    raise ValueError(\n        'top_level_all_gather works when all mesh axes of context mesh are'\n        f' explicit. Got {get_abstract_mesh()}')\n  x_flat, treedef = tree_flatten(xs)\n  out_sharding_flat = api_util.flatten_axis_resources(\n      \"top_level_all_gather out_sharding\", treedef, out_sharding,\n      tupled_args=True)\n  x_avals_flat = [core.typeof(x) for x in x_flat]\n  out_flat = [_top_level_ag(x, aval, sh, multi_dim)\n              for x, aval, sh in zip(x_flat, x_avals_flat, out_sharding_flat)]\n  return tree_unflatten(treedef, out_flat)\n","sourceCodeStart":2230,"sourceCodeEnd":2259,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L2230-L2259","documentation":"top_level_all_gather only supports context meshes whose axes are all explicit (created with axis_types where every axis is named/explicit). If the ambient abstract mesh has positional/implicit axes, the API refuses to run.","triggerScenarios":"Calling top_level_all_gather while the context mesh (set via mesh context manager or with_mesh) contains non-explicit axes (axis_types=None entries).","commonSituations":"Using new-style jax.make_mesh(..., axis_types=...) partially explicit, or forgetting to enter an explicit mesh context before the call.","solutions":["Enter an explicit mesh context before calling (e.g. `with mesh: top_level_all_gather(...)`), where mesh was built with all-explicit axes","Rebuild the mesh with jax.make_mesh(shape, names, axis_types=(...)) with no None entries","Or use shard_map/pjit-based all_gather instead of top_level_all_gather"],"exampleFix":"# before\ntop_level_all_gather(x, out_sh)  # context mesh has implicit axes\n\n# after\nwith explicit_mesh:\n  top_level_all_gather(x, out_sh)","handlingStrategy":"validation","validationCode":"from jax.experimental.mesh_utils import get_abstract_mesh  # or jax._src.mesh\nm = get_abstract_mesh()\nassert m.are_all_axes_explicit, f'need all-explicit mesh, got {m}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build meshes with make_mesh(..., axis_types=(...)) with no None entries","Enter the mesh context (`with mesh:`) around top-level collective calls"],"tags":["jax","all-gather","explicit-mesh","context-mesh"],"backgroundTag":"jax-mesh-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}