{"record":{"id":"f2a0215d72ea3c7c","repo":"jax-ml/jax","slug":"out-sharding-passed-to-top-level-all-gather-cannot","errorCode":null,"errorMessage":"out_sharding passed to top_level_all_gather cannot be {out_sh_}. It should be a PartitionSpec or NamedSharding.","messagePattern":"out_sharding passed to top_level_all_gather cannot be (.+?)\\. It should be a PartitionSpec or NamedSharding\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":2203,"sourceCode":"      newly_manual_axes=newly_manual_axes, debug_info=debug_info,\n      check_vma=check_vma)\n  out_vals, ref_vals = split_list(out_and_ref_vals, [len(jaxpr.outvars)])\n  ref_vals_ = iter(ref_vals)\n  new_invals = [next(ref_vals_) if isinstance(a, AbstractRef) else None\n                for a in ctx.in_avals]\n  assert next(ref_vals_, None) is None\n  return new_invals, out_vals\n\ndef _repspec(aval):\n  return aval.nospec(empty_abstract_mesh, False, ())\n\n# ----------------------- top level collectives --------------------------------\n\ndef _top_level_ag(x, aval, out_sh_, multi_dim):\n  assert aval.sharding.mesh.are_all_axes_explicit, aval.sharding.mesh\n  out_sh = canonicalize_sharding(out_sh_, \"top_level_all_gather\")\n  if out_sh is None:\n    raise ValueError(\n        f'out_sharding passed to top_level_all_gather cannot be {out_sh_}. It'\n        ' should be a PartitionSpec or NamedSharding.')\n  if aval.sharding.mesh != out_sh.mesh:\n    raise ValueError(\n        f'Input sharding mesh {aval.sharding.mesh} should be equal to'\n        f' out_sharding mesh {out_sh.mesh}')\n\n  in_spec = aval.sharding.spec\n  out_spec = out_sh.spec._normalized_spec_for_aval(len(in_spec))\n  if config.remove_size_one_mesh_axis_from_type.value:\n    out_spec = remove_size_one_mesh_axis_from_spec(out_spec, out_sh.mesh)\n\n  def f_shmap(x):\n    # Maybe this can just be 1 AG where we gather in a new dim and then do\n    # AG(new_dim) -> reshape -> transpose -> reshape but it might be expensive.\n    count = 0\n    for axis, (i, o) in enumerate(zip(in_spec.partitions, out_spec.partitions)):\n      if i == o:","sourceCodeStart":2185,"sourceCodeEnd":2221,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L2185-L2221","documentation":"top_level_all_gather canonicalizes its out_sharding argument; if canonicalization returns None the value is not a PartitionSpec or NamedSharding (e.g. a raw string, GSPMDSharding, or wrong object), and JAX raises this error.","triggerScenarios":"Calling jax.experimental.top_level_all_gather (shard_map module) with out_sharding that is not a PartitionSpec/NamedSharding, such as a string spec or an OpSharding proto.","commonSituations":"Passing a PjitSharding/GSPMDSharding or hand-built object where a user-facing PartitionSpec is expected; API confusion with pjit's in_shardings which accept more types.","solutions":["Pass a NamedSharding(mesh, PartitionSpec(...)) or a plain PartitionSpec as out_sharding","If you have an HloSharding, convert it to a NamedSharding first","Check for None/typo'd variables being passed through"],"exampleFix":"# before\ntop_level_all_gather(x, out_sharding='data')\n\n# after\nfrom jax.sharding import NamedSharding, PartitionSpec as P\ntop_level_all_gather(x, out_sharding=NamedSharding(mesh, P('data')))","handlingStrategy":"type-guard","validationCode":"from jax.sharding import PartitionSpec, NamedSharding\nassert isinstance(out_sharding, (PartitionSpec, NamedSharding)), f'bad out_sharding type: {type(out_sharding)}'","typeGuard":"def is_valid_out_sharding(s) -> bool:\n    from jax.sharding import PartitionSpec, NamedSharding\n    return isinstance(s, (PartitionSpec, NamedSharding))","tryCatchPattern":null,"preventionTips":["Construct out_sharding via NamedSharding(mesh, P(...)) at the call site","Never pass HloSharding/OpSharding protos to top_level_all_gather"],"tags":["jax","shard-map","all-gather","sharding-type"],"backgroundTag":"invalid-sharding-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}