{"record":{"id":"4a2facb33c6bf21a","repo":"jax-ml/jax","slug":"unbound-axis-name-axis-name-4a2fac","errorCode":null,"errorMessage":"unbound axis name: {axis_name}","messagePattern":"unbound axis name: (.+?)","errorType":"validation","errorClass":"NameError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/pxla.py","lineNumber":450,"sourceCode":"\n    if self.mut is None:\n      return out\n    else:\n      out_ = []\n      for i, o in zip(self.mut.out_mut, out):\n        if i is not None:\n          try: args[i]._refs._buf._replace_with(o)\n          except AttributeError: pass  # TODO(mattjj): remove float0\n        else:\n          out_.append(o)\n      return out_\n\n\ndef _axis_read(axis_names, axis_name):\n  try:\n    return max(i for i, name in enumerate(axis_names) if name == axis_name)\n  except ValueError:\n    raise NameError(f\"unbound axis name: {axis_name}\") from None\n\ndef axis_groups(axis_ctx, name) -> tuple[tuple[int, ...]]:\n  assert not isinstance(axis_ctx, sharding_impls.ShardingContext)\n  size = axis_ctx.mesh.size\n  axis_names = axis_ctx.mesh.axis_names\n  axis_sizes = axis_ctx.mesh.axis_sizes\n  if not isinstance(name, (list, tuple)):\n    name = (name,)\n  mesh_axes = tuple(unsafe_map(partial(_axis_read, axis_names), name))\n  trailing_size, ragged = divmod(size, math.prod(axis_sizes))\n  assert not ragged\n  mesh_spec = axis_sizes + (trailing_size,)\n  return _axis_groups(mesh_spec, mesh_axes)\n\ndef _axis_groups(mesh_spec, mesh_axes):\n  \"\"\"Computes replica group ids for a collective performed over a subset of the mesh.\n\n  Args:","sourceCodeStart":432,"sourceCodeEnd":468,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/pxla.py#L432-L468","documentation":"Inside sharded computations with a Mesh, collective operations reference axes by name (e.g. axis_index('data'), psum(..., axis_name='data')). _axis_read looks up the name among the context mesh's axis names and raises NameError if absent. The `max(...)` makes repeated axis names use the last (innermost) occurrence.","triggerScenarios":"Using a collective with an axis_name not present in the active mesh: psum(x, 'model') under a Mesh with only a 'data' axis; using xmap/pjit with axis names mismatched between the annotation and the mesh; typos or renaming mesh axes without updating collectives.","commonSituations":"Renaming mesh axes; copy-pasting collective code between models with different mesh layouts; nested/sharded_jit contexts where the axis context isn't the expected mesh.","solutions":["Use an axis name that exists in the Mesh: check mesh.axis_names and pass one of those to the collective","Fix typos/differences between the declared Mesh axis names and the axis_name argument","If the axis genuinely should exist, verify the collective is evaluated under the right mesh context (e.g. within the sharded computation, not outside)"],"exampleFix":"# before\nmesh = Mesh(devices, axis_names=('data',))\n... psum(x, axis_name='model')\n\n# after\nmesh = Mesh(devices, axis_names=('data',))\n... psum(x, axis_name='data')","handlingStrategy":"type-guard","validationCode":"def check_axis(mesh, axis_name):\n    assert axis_name in mesh.axis_names, \\\n        f'{axis_name!r} not in mesh axes {mesh.axis_names}'","typeGuard":"def axis_is_bound(mesh, axis_name: str) -> bool:\n    return axis_name in mesh.axis_names","tryCatchPattern":null,"preventionTips":["Derive collective axis_name arguments from mesh.axis_names programmatically","Rename mesh axes atomically: update all collectives in the same change"],"tags":["jax","mesh","collectives","axis-name","sharding"],"backgroundTag":"name-not-found","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}