jax-ml/jax · error · TypeError

smap in_axes must be an int, None, jax.sharding.Infer, or a

Error message

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}.

What it means

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.

Source

Thrown at jax/_src/shard_map.py:226

  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,
                       is_leaf=lambda x: x is None))
  out_specs = tree_map(partial(_axes_to_pspec, axis_name), out_axes,
                       is_leaf=lambda x: x is None)
  return _shard_map(f, mesh=None, in_specs=in_specs, out_specs=out_specs,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use integers: in_axes=0 or in_axes=(0, None)
2. If you want to specify sharding by mesh axis names, use jax.shard_map with in_specs=P('i', 'j') instead of smap
3. Verify tuple length matches the number of positional arguments f accepts

Example fix

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

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

def check_smap_in_axes(in_axes, n_args):
    ok = in_axes is None or in_axes is Infer or isinstance(in_axes, int) or (
        isinstance(in_axes, tuple) and len(in_axes) == n_args)
    if not ok:
        raise TypeError(f'bad in_axes: {in_axes!r}')

def is_valid_in_axes(in_axes) -> bool:
    return in_axes is None or isinstance(in_axes, (int, tuple))

When it happens

Trigger: 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.

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

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 axis_name should be a `str` or a `Hashable`, but got {a
• smap in_axes must be an int, None, jax.sharding.Infer, or (n
• 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

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