{"record":{"id":"8a834bd1f2e624d0","repo":"jax-ml/jax","slug":"unbound-axis-name-axis-name","errorCode":null,"errorMessage":"unbound axis name: {axis_name}","messagePattern":"unbound axis name: (.+?)","errorType":"exception","errorClass":"NameError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":1376,"sourceCode":"  spmd_axis_names : frozenset[AxisName]\n  explicit_mesh_axis_names: frozenset[AxisName]\n\n  @staticmethod\n  @weak_value_interner\n  def _create(axis_sizes, spmd_axis_names, explicit_mesh_axis_names):\n    obj = object.__new__(AxisEnv)\n    object.__setattr__(obj, 'axis_sizes', axis_sizes)\n    object.__setattr__(obj, 'spmd_axis_names', spmd_axis_names)\n    object.__setattr__(obj, 'explicit_mesh_axis_names', explicit_mesh_axis_names)\n    return obj\n\n  def __new__(cls, axis_sizes=FrozenDict({}), spmd_axis_names=frozenset(),\n              explicit_mesh_axis_names=frozenset()):\n    return cls._create(axis_sizes, spmd_axis_names, explicit_mesh_axis_names)\n\n  def axis_size(self, axis_name):\n    if axis_name not in self.axis_sizes:\n      raise NameError(f\"unbound axis name: {axis_name}\")\n    else:\n      return self.axis_sizes[axis_name]\n\n  def axis_exists(self, axis_name):\n    return axis_name in self.axis_sizes\n\n  def axis_names(self):\n    return tuple(k for k in self.axis_sizes)\n\n  def pop_pure(self, axis_name):\n    new_sizes = dict(self.axis_sizes)\n    new_sizes.pop(axis_name)\n    return AxisEnv(FrozenDict(new_sizes), self.spmd_axis_names,\n                   self.explicit_mesh_axis_names)\n\n  def extend_pure(self, name_size_pairs):\n    new_sizes = dict(self.axis_sizes)\n    new_sizes.update((name, size) for name, size in name_size_pairs","sourceCodeStart":1358,"sourceCodeEnd":1394,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L1358-L1394","documentation":"JAX NameError thrown by AxisEnv.axis_size when a named axis (from pmap/xmap/sharding annotations) is looked up but is not bound in the current axis environment. JAX tracks named logical axes explicitly, and referencing one that was never introduced into the trace's axis environment is an error.","triggerScenarios":"Calling axis_size(axis_name) on an AxisEnv lacking that name; using vmap/xmap/pmap with an out_axes/in_axes axis_name not declared; using an axis name in collectives (e.g., psum, all_gather) that is not among the mapped axis names of the enclosing transformation.","commonSituations":"Typos in axis_name strings; using psum(..., axis_name='i') inside vmap without axis_name='i'; refactoring renamed axes in one place but not the collective call; mixing pmap axis_name with vmap-only contexts.","solutions":["Check the spelling/case of axis_name against the name passed to the enclosing vmap/pmap/xmap","Ensure the collective (psum/pmax/all_reduce...) is called inside the transformation that declared that axis name","Nest transformations so that every axis_name referenced is declared by an enclosing vmap(name=...) or pmap(axis_name=...)","Print/inspect the AxisEnv (axis_sizes keys) in a minimal repro to see which names are bound"],"exampleFix":"// before\nout = jax.vmap(fn)(x)  # fn uses lax.psum(x, 'batch')\n\n// after\nout = jax.vmap(fn, axis_name='batch')(x)","handlingStrategy":"validation","validationCode":"def check_axis_bound(axis_name, declared_names):\n    assert axis_name in declared_names, f\"axis {axis_name!r} not in {declared_names}\"","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Declare axis names centrally as constants and reference them everywhere","Always pair psum/all_gather with the same axis_name string used in vmap/pmap","Test collectives in a minimal vmap repro before composing transformations"],"tags":["jax","axis-name","vmap","pmap","collectives"],"backgroundTag":"jax-unbound-axis-name","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}