{"record":{"id":"f97a4ae3c23d486f","repo":"jax-ml/jax","slug":"smap-axis-name-should-be-a-str-or-a-hashable","errorCode":null,"errorMessage":"smap axis_name should be a `str` or a `Hashable`, but got {axis_name}","messagePattern":"smap axis_name should be a `str` or a `Hashable`, but got (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":222,"sourceCode":"      (tuple/list/dict) thereof indicating where the mapped axis should appear\n      in the output.\n    axis_name: ``mesh`` axis name over which the function ``f`` is manual.\n\n  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,","sourceCodeStart":204,"sourceCodeEnd":240,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L204-L240","documentation":"jax.experimental.shard_map.smap requires axis_name to be a single hashable (typically a str), not a list or tuple. Passing multiple names as a sequence is explicitly rejected to catch confusion with shard_map's axis_names set API.","triggerScenarios":"Calling jax.experimental.smap(f, mesh=..., axis_name=('i','j'), ...) or axis_name=['i'] — any list/tuple value for axis_name.","commonSituations":"Users migrating from lax.pmap's axis_name (which accepts single names but where people habitually pass collections), or from jax.lax.map / vmap-style APIs where axis arguments are tuples.","solutions":["Pass a single string/hashable: axis_name='i'","If multiple axes are needed, use jax.shard_map directly with axis_names=frozenset({'i','j'})"],"exampleFix":"// before\njax.experimental.smap(f, mesh=mesh, in_axes=0, out_axes=0, axis_name=('i','j'))\n\n// after\njax.shard_map(f, mesh=mesh, in_specs=P('i','j'), out_specs=P('i','j'), axis_names=frozenset({'i','j'}))","handlingStrategy":"type-guard","validationCode":"from typing import Hashable\n\ndef valid_axis_name(name):\n    return isinstance(name, str) or (isinstance(name, Hashable) and not isinstance(name, (list, tuple)))","typeGuard":"def is_valid_smap_axis_name(name) -> bool:\n    return not isinstance(name, (list, tuple)) and isinstance(name, Hashable)","tryCatchPattern":null,"preventionTips":["smap takes exactly one axis_name; use jax.shard_map with a set for multiple axes","Linter-annotate axis_name as str in your wrappers"],"tags":["jax","shard-map","smap","axis-name","distributed"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}