{"record":{"id":"d356e841eb5f1f6f","repo":"jax-ml/jax","slug":"smap-out-axes-must-be-an-int-none-or-nested-co","errorCode":null,"errorMessage":"smap out_axes must be an int, None, or (nested) container with those types as leaves, but got {out_axes}.","messagePattern":"smap out_axes must be an int, None, 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":236,"sourceCode":"def _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,\n    in_specs: Specs, out_specs: Specs, axis_names: Set[AxisName],\n    check_vma: bool, _smap: bool = False) -> F:\n  if not callable(f):\n    raise TypeError(\"shard_map requires a callable for its first argument, \"","sourceCodeStart":218,"sourceCodeEnd":254,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L218-L254","documentation":"jax.experimental.smap requires out_axes to be an int, None, or a (nested) pytree container whose leaves are ints. Non-integer leaves (strings, floats, lists) in out_axes trigger this TypeError.","triggerScenarios":"Calling smap with out_axes='i', out_axes=[0], or out_axes=(0, 'j').","commonSituations":"Symmetric confusion with in_axes: users specify mesh axis names or Python lists where integer output axes are expected; also copied-over PartitionSpec values from shard_map code.","solutions":["Use integers or None: out_axes=0 or out_axes=(0, None)","Switch to jax.shard_map with out_specs=P(...) for name-based output sharding"],"exampleFix":"// before\njax.experimental.smap(f, mesh=mesh, in_axes=0, out_axes='i', axis_name='i')\n\n// after\njax.experimental.smap(f, mesh=mesh, in_axes=0, out_axes=0, axis_name='i')","handlingStrategy":"type-guard","validationCode":"from jax.tree_util import tree_leaves\nassert all(isinstance(l, int) for l in tree_leaves(out_axes)), out_axes","typeGuard":"def is_valid_out_axes(spec) -> bool:\n    return all(isinstance(l, int) for l in tree_leaves(spec))","tryCatchPattern":null,"preventionTips":["out_axes mirrors in_axes: ints only","Use shard_map out_specs=P(...) for name-based output sharding"],"tags":["jax","smap","out-axes","validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}