{"record":{"id":"dde95a629b5071ed","repo":"jax-ml/jax","slug":"smap-in-axes-must-be-an-int-none-jax-sharding-in","errorCode":null,"errorMessage":"smap in_axes must be an int, None, jax.sharding.Infer, or a tuple of entries corresponding to the positional arguments passed to the function, but got {in_axes}.","messagePattern":"smap in_axes must be an int, None, jax\\.sharding\\.Infer, or a tuple of entries corresponding to the positional arguments passed to the function, but got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":226,"sourceCode":"  Returns:\n    A callable representing a mapped version of ``f``, which accepts positional\n    arguments corresponding to those of ``f`` and produces output corresponding\n    to that of ``f``.\n  \"\"\"\n  kwargs = dict(in_axes=in_axes, out_axes=out_axes, axis_name=axis_name)\n  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,","sourceCodeStart":208,"sourceCodeEnd":244,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L208-L244","documentation":"jax.experimental.shard_map.smap validates in_axes at call time: it must be None, an int, jax.sharding.Infer, or a tuple whose entries correspond to the function's positional arguments. Any other type (str, dict, list, float, ...) triggers this TypeError.","triggerScenarios":"Calling smap with in_axes='i' (a PartitionSpec-style name), in_axes=[0, 1] (list instead of tuple), in_axes=0.5, or a tuple of wrong length relative to the function's positional args.","commonSituations":"Confusion between smap's vmap-like in_axes (integer axis indices) and shard_map's in_specs (PartitionSpec names). Users often pass a mesh axis name string where a positional axis integer is required.","solutions":["Use integers: in_axes=0 or in_axes=(0, None)","If you want to specify sharding by mesh axis names, use jax.shard_map with in_specs=P('i', 'j') instead of smap","Verify tuple length matches the number of positional arguments f accepts"],"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":"def check_smap_in_axes(in_axes, n_args):\n    ok = in_axes is None or in_axes is Infer or isinstance(in_axes, int) or (\n        isinstance(in_axes, tuple) and len(in_axes) == n_args)\n    if not ok:\n        raise TypeError(f'bad in_axes: {in_axes!r}')","typeGuard":"def is_valid_in_axes(in_axes) -> bool:\n    return in_axes is None or isinstance(in_axes, (int, tuple))","tryCatchPattern":null,"preventionTips":["Remember smap axes are vmap-style ints, not mesh names","Keep tuple length equal to positional arg count","Use jax.shard_map + PartitionSpec for name-based sharding"],"tags":["jax","smap","shard-map","in-axes","distributed"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}