{"record":{"id":"583b6acbdedd5e45","repo":"jax-ml/jax","slug":"smap-in-axes-must-be-an-int-none-jax-sharding-in-583b6a","errorCode":null,"errorMessage":"smap in_axes must be an int, None, jax.sharding.Infer, or (nested) container with those types as leaves, but got {in_axes}.","messagePattern":"smap in_axes must be an int, None, jax\\.sharding\\.Infer, or \\(nested\\) container with those types as leaves, but got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":232,"sourceCode":"  if f is None:\n    return lambda g: _smap(g, **kwargs)\n  return _smap(f, **kwargs)\n\ndef _smap[F: Callable](\n    f: F, *, in_axes: int | None | InferFromArgs | tuple[Any, ...],\n    out_axes: Any, axis_name: AxisName) -> F:\n  if isinstance(axis_name, (list, tuple)):\n    raise TypeError(\n        f\"smap axis_name should be a `str` or a `Hashable`, but got {axis_name}\")\n  if (in_axes is not None and in_axes is not Infer and\n      not isinstance(in_axes, (int, tuple))):\n    raise TypeError(\n        \"smap in_axes must be an int, None, jax.sharding.Infer, or a tuple of\"\n        \" entries corresponding to the positional arguments passed to the\"\n        f\" function, but got {in_axes}.\")\n  if (in_axes is not Infer and\n      not all(isinstance(l, int) for l in tree_leaves(in_axes))):\n    raise TypeError(\n        \"smap in_axes must be an int, None, jax.sharding.Infer, or (nested)\"\n        f\" container with those types as leaves, but got {in_axes}.\")\n  if not all(isinstance(l, int) for l in tree_leaves(out_axes)):\n    raise TypeError(\"smap out_axes must be an int, None, or (nested) container \"\n                    f\"with those types as leaves, but got {out_axes}.\")\n\n  in_specs = (Infer if in_axes is Infer else\n              tree_map(partial(_axes_to_pspec, axis_name), in_axes,\n                       is_leaf=lambda x: x is None))\n  out_specs = tree_map(partial(_axes_to_pspec, axis_name), out_axes,\n                       is_leaf=lambda x: x is None)\n  return _shard_map(f, mesh=None, in_specs=in_specs, out_specs=out_specs,\n                    axis_names={axis_name}, check_vma=True, _smap=True)\n\n\n@partial(traceback_util.api_boundary, repro_api_name=\"jax.shard_map\")\ndef _shard_map[F: Callable](\n    f: F, *, mesh: Mesh | AbstractMesh | None,","sourceCodeStart":214,"sourceCodeEnd":250,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L214-L250","documentation":"For jax.experimental.smap, when in_axes is a (possibly nested) pytree container, every leaf must be an int (or None entries as leaves via is_leaf). This error fires when leaves are non-integer, e.g. strings or floats, inside a nested tuple/tree structure.","triggerScenarios":"Calling smap with in_axes=((0, 'i'),) or nested containers containing mesh-axis-name strings instead of integer positional axes.","commonSituations":"Users copying shard_map PartitionSpec patterns into smap's vmap-style axes argument, or nesting tuples per pytree structure of the inputs but using names instead of ints.","solutions":["Replace every leaf with an integer axis or None: in_axes=((0, None),)","Use jax.shard_map with in_specs=P(...) if you need name-based sharding"],"exampleFix":"// before\njax.experimental.smap(f, mesh=mesh, in_axes=(('i', None),), out_axes=0, axis_name='i')\n\n// after\njax.experimental.smap(f, mesh=mesh, in_axes=((0, None),), out_axes=0, axis_name='i')","handlingStrategy":"type-guard","validationCode":"from jax.tree_util import tree_leaves\n\ndef leaves_are_ints(spec):\n    return all(l is None or isinstance(l, int) for l in tree_leaves(spec))","typeGuard":"def is_valid_axes_tree(spec) -> bool:\n    return all(isinstance(l, int) for l in tree_leaves(spec))","tryCatchPattern":null,"preventionTips":["Never put mesh axis names inside smap in_axes","Validate leaves with tree_leaves before calling in dynamic code"],"tags":["jax","smap","in-axes","pytree","validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}