{"record":{"id":"c97cebf4d6b65cc7","repo":"jax-ml/jax","slug":"found-an-unbound-axis-name-name-to-fix-this-p","errorCode":null,"errorMessage":"Found an unbound axis name: {name}. To fix this, please call {api_name} under `jax.shard_map`.","messagePattern":"Found an unbound axis name: (.+?)\\. To fix this, please call (.+?) under `jax\\.shard_map`\\.","errorType":"exception","errorClass":"NameError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1033,"sourceCode":"  check_unreduced_args([aval], axes, 'psum')\n  out_aval = ShapedArray(\n      lax._reduce_op_shape_rule(aval, axes=pos_axes), aval.dtype,\n      sharding=lax._reduce_op_sharding_rule(aval, axes=pos_axes))\n  return out_aval, {core.NamedAxisEffect(axis) for axis in named_axes}\n\n# TODO(yashkatariya): Replace this with _psum_invariant_abstract_eval\ndef _pmin_pmax_abstract_eval(name, aval, *, axes, axis_index_groups):\n  if not config._check_vma.value:\n    return _allreduce_effectful_abstract_eval(\n        aval, axes=axes, axis_index_groups=axis_index_groups)\n  return _psum_invariant_abstract_eval(name, aval, axes=axes)\n\ndef _check_axis_names(axes, api_name):\n  named_axes = tuple(axis for axis in axes if not isinstance(axis, int))\n  axis_env = core.get_axis_env()\n  for name in named_axes:\n    if not axis_env.axis_exists(name):\n      raise NameError(\n          f\"Found an unbound axis name: {name}. To fix this, please call\"\n          f\" {api_name} under `jax.shard_map`.\")\n\ndef _lower_reducer_into_block(ctx, prim, scalar_aval, block):\n  with ir.InsertionPoint(block):\n    lower_reducer = mlir.lower_fun(prim.bind, multiple_results=False)\n    reducer_ctx = ctx.replace(\n        primitive=None, avals_in=[scalar_aval] * 2, avals_out=[scalar_aval]\n    )\n    out_nodes = lower_reducer(reducer_ctx, *block.arguments)\n    flat_out_nodes, _ = mlir.ir_tree_registry.flatten(out_nodes)\n    if isinstance(block.owner, func_dialect.FuncOp):\n      func_dialect.return_(flat_out_nodes)\n    else:\n      hlo.return_(flat_out_nodes)\n\n\ndef _build_reducer_func_op(ctx, prim, aval_in):","sourceCodeStart":1015,"sourceCodeEnd":1051,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1015-L1051","documentation":"Collectives like psum, all_gather, ppermute, all_to_all require a bound mesh axis name. _check_axis_names looks the name up in the current axis environment and raises NameError if no vmap/shard_map frame introduced it.","triggerScenarios":"Calling jax.lax.psum(x, 'i') (or all_to_all/all_gather/psend) outside jax.shard_map or without the axis being declared on the mesh/shard_map in_axes.","commonSituations":"Running distributed code eagerly (no shard_map) during debugging; axis name typo ('dp' vs 'data'); mesh declared but function not wrapped in shard_map.","solutions":["Wrap the call in jax.shard_map with a mesh that declares the axis name","Check for a typo in the axis name against Mesh(axis_names=...)","If using vmap-only, wrap with jax.vmap(..., axis_name='i')"],"exampleFix":"// before\ny = jax.lax.psum(x, 'i')\n// after\ny = jax.shard_map(lambda b: jax.lax.psum(b, 'i'), mesh)(x)","handlingStrategy":"validation","validationCode":"def assert_axis_bound(mesh, name):\n    assert name in mesh.axis_names, f'{name} not in mesh axes {mesh.axis_names}'\n\n# and always call collectives inside shard_map","typeGuard":null,"tryCatchPattern":"try:\n    y = jax.shard_map(f, mesh)(x)\nexcept NameError as e:\n    if 'unbound axis name' in str(e): raise RuntimeError(f'mesh misconfigured: {e}')\n    raise","preventionTips":["Declare and assert axis names against mesh.axis_names","Always wrap collective calls in shard_map or vmap(axis_name=...)"],"tags":["jax","mesh","shard-map","unbound-axis"],"backgroundTag":"unbound-name","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}