jax-ml/jax · error · TypeError

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

Error message

smap in_axes must be an int, None, jax.sharding.Infer, or (nested) container with those types as leaves, but got {in_axes}.

What it means

For jax.experimental.smap, when in_axes is a (possibly nested) pytree container, every leaf must be an int (or None entries as leaves via is_leaf). This error fires when leaves are non-integer, e.g. strings or floats, inside a nested tuple/tree structure.

Source

Thrown at jax/_src/shard_map.py:232

  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,
                    axis_names={axis_name}, check_vma=True, _smap=True)


@partial(traceback_util.api_boundary, repro_api_name="jax.shard_map")
def _shard_map[F: Callable](
    f: F, *, mesh: Mesh | AbstractMesh | None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace every leaf with an integer axis or None: in_axes=((0, None),)
2. Use jax.shard_map with in_specs=P(...) if you need name-based sharding

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')

from jax.tree_util import tree_leaves

def leaves_are_ints(spec):
    return all(l is None or isinstance(l, int) for l in tree_leaves(spec))

def is_valid_axes_tree(spec) -> bool:
    return all(isinstance(l, int) for l in tree_leaves(spec))

When it happens

Trigger: Calling smap with in_axes=((0, 'i'),) or nested containers containing mesh-axis-name strings instead of integer positional axes.

Common situations: Users copying shard_map PartitionSpec patterns into smap's vmap-style axes argument, or nesting tuples per pytree structure of the inputs but using names instead of ints.

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

• {full_name} must be a pytree prefix with bool leaves or a tu
• smap in_axes must be an int, None, jax.sharding.Infer, or a
• smap out_axes must be an int, None, or (nested) container wi
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64

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