{"record":{"id":"a6cc504dc3d8c82e","repo":"jax-ml/jax","slug":"in-axes-was-not-specified-when-axis-name-axis-nam","errorCode":null,"errorMessage":"in_axes was not specified when axis_name={axis_names} was of type {axis_types} / shard_map in_specs argument must be a pytree of `jax.sharding.PartitionSpec` instances, but it was `None` when {axis_names=} are of type {axis_types}","messagePattern":"in_axes was not specified when axis_name=(.+?) was of type (.+?) / shard_map in_specs argument must be a pytree of `jax\\.sharding\\.PartitionSpec` instances, but it was `None` when (.+?) are of type (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":412,"sourceCode":"    axis_names = frozenset(axis_names)\n  if not axis_names:\n    axis_names = mesh_axis_names_wo_vmap\n  if not axis_names.issubset(mesh_axis_names_wo_vmap):\n    raise ValueError(\n        f\"jax.shard_map requires axis_names={axis_names} to be a subset of \"\n        f\"mesh.axis_names={mesh_axis_names_wo_vmap}\")\n\n  if (in_specs is Infer and\n      not all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):\n    axis_types = ', '.join(str(mesh._name_to_type[a]) for a in axis_names)\n    if _smap:\n      msg = (f\"in_axes was not specified when axis_name={axis_names} was of\"\n             f\" type {axis_types}\")\n    else:\n      msg = (\"shard_map in_specs argument must be a pytree of\"\n             \" `jax.sharding.PartitionSpec` instances, but it was `None` when\"\n             f\" {axis_names=} are of type {axis_types}\")\n    raise TypeError(msg)\n\n  if in_specs is not Infer and in_specs is not None:\n    _check_specs(SpecErrorType.input, in_specs, axis_names)\n    _check_unreduced(SpecErrorType.input, mesh, axis_names, in_specs)\n  _check_specs(SpecErrorType.out, out_specs, axis_names)\n  _check_unreduced(SpecErrorType.out, mesh, axis_names, out_specs)\n  return mesh, axis_names\n\n\ndef _manual_spec(manual_axes, spec: P, mesh) -> P:\n  out: list[str | tuple[str | None, ...] | None] = []\n  s: str | None | tuple[str, ...]\n  for s in spec.partitions:\n    if s is None:\n      out.append(s)\n    elif isinstance(s, tuple):\n      temp = [p if p in manual_axes else None for p in s]\n      while temp and temp[-1] is None:","sourceCodeStart":394,"sourceCodeEnd":430,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L394-L430","documentation":"shard_map requires either explicit PartitionSpec in_specs/in_axes or an explicit-mesh context when axis_names include non-Explicit axis types (e.g. autodiff/vmap-introduced axes). The message has two variants: for smap it says in_axes was unspecified; for shard_map it says in_specs was None while axis_names were of a non-Explicit type. In short, you can't infer specs when the named axes aren't concrete mesh axes.","triggerScenarios":"Calling smap without in_axes while axis_name refers to a non-Explicit axis type, or calling shard_map with in_specs=None when any named axis is not an Explicit mesh axis (e.g. inside vmap/grad where implicit axes exist).","commonSituations":"Composing vmap or grad with shard_map/smap; using abstract/multi-stage meshes where axes have types like Autodiff or Replicated; migrating code that relied on inference that now fails under typed mesh axes.","solutions":["Specify in_axes/in_specs explicitly so JAX needn't infer specs from non-Explicit axes","Restrict axis_names to Explicit mesh axes, or mark the mesh axes as explicit when constructing the abstract mesh"],"exampleFix":"// before\njax.experimental.smap(f, mesh=mesh, out_axes=0, axis_name='i')  # 'i' not Explicit, no in_axes\n\n// after\njax.experimental.smap(f, mesh=mesh, in_axes=0, out_axes=0, axis_name='i')","handlingStrategy":"validation","validationCode":"# always specify in_axes/in_specs when non-Explicit axes may be present\nif any(mesh._name_to_type[a] != AxisType.Explicit for a in axis_names):\n    assert in_axes is not None and in_axes is not Infer, 'specify in_axes explicitly'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pass in_axes=0 / in_specs=P(...) explicitly in vmap+shard_map code","Prefer Explicit mesh axes for shard_map names"],"tags":["jax","shard-map","smap","axis-types","in-specs"],"backgroundTag":"missing-sharding-spec","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}