jax-ml/jax · error · TypeError

smap axis_name should be a `str` or a `Hashable`, but got {a

Error message

smap axis_name should be a `str` or a `Hashable`, but got {axis_name}

What it means

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.

Source

Thrown at jax/_src/shard_map.py:222

      (tuple/list/dict) thereof indicating where the mapped axis should appear
      in the output.
    axis_name: ``mesh`` axis name over which the function ``f`` is manual.

  Returns:
    A callable representing a mapped version of ``f``, which accepts positional
    arguments corresponding to those of ``f`` and produces output corresponding
    to that of ``f``.
  """
  kwargs = dict(in_axes=in_axes, out_axes=out_axes, axis_name=axis_name)
  if f is None:
    return lambda g: _smap(g, **kwargs)
  return _smap(f, **kwargs)

def _smap[F: Callable](
    f: F, *, in_axes: int | None | InferFromArgs | tuple[Any, ...],
    out_axes: Any, axis_name: AxisName) -> F:
  if isinstance(axis_name, (list, tuple)):
    raise TypeError(
        f"smap axis_name should be a `str` or a `Hashable`, but got {axis_name}")
  if (in_axes is not None and in_axes is not Infer and
      not isinstance(in_axes, (int, tuple))):
    raise TypeError(
        "smap in_axes must be an int, None, jax.sharding.Infer, or a tuple of"
        " entries corresponding to the positional arguments passed to the"
        f" function, but got {in_axes}.")
  if (in_axes is not Infer and
      not all(isinstance(l, int) for l in tree_leaves(in_axes))):
    raise TypeError(
        "smap in_axes must be an int, None, jax.sharding.Infer, or (nested)"
        f" container with those types as leaves, but got {in_axes}.")
  if not all(isinstance(l, int) for l in tree_leaves(out_axes)):
    raise TypeError("smap out_axes must be an int, None, or (nested) container "
                    f"with those types as leaves, but got {out_axes}.")

  in_specs = (Infer if in_axes is Infer else
              tree_map(partial(_axes_to_pspec, axis_name), in_axes,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a single string/hashable: axis_name='i'
2. If multiple axes are needed, use jax.shard_map directly with axis_names=frozenset({'i','j'})

Example fix

// before
jax.experimental.smap(f, mesh=mesh, in_axes=0, out_axes=0, axis_name=('i','j'))

// after
jax.shard_map(f, mesh=mesh, in_specs=P('i','j'), out_specs=P('i','j'), axis_names=frozenset({'i','j'}))

from typing import Hashable

def valid_axis_name(name):
    return isinstance(name, str) or (isinstance(name, Hashable) and not isinstance(name, (list, tuple)))

def is_valid_smap_axis_name(name) -> bool:
    return not isinstance(name, (list, tuple)) and isinstance(name, Hashable)

When it happens

Trigger: Calling jax.experimental.smap(f, mesh=..., axis_name=('i','j'), ...) or axis_name=['i'] — any list/tuple value for axis_name.

Common situations: 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.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• smap in_axes must be an int, None, jax.sharding.Infer, or a
• in_specs passed to shard_map: {s} does not match the specs o
• shard_map out_specs rank error (msg from _spec_rank_error)
• The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t
• The context mesh {ctx_mesh} should match the mesh passed to

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f97a4ae3c23d486f. Report an issue: GitHub.