jax-ml/jax · error · TypeError

smap out_axes must be an int, None, or (nested) container wi

Error message

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

What it means

jax.experimental.smap requires out_axes to be an int, None, or a (nested) pytree container whose leaves are ints. Non-integer leaves (strings, floats, lists) in out_axes trigger this TypeError.

Source

Thrown at jax/_src/shard_map.py:236

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,
    in_specs: Specs, out_specs: Specs, axis_names: Set[AxisName],
    check_vma: bool, _smap: bool = False) -> F:
  if not callable(f):
    raise TypeError("shard_map requires a callable for its first argument, "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use integers or None: out_axes=0 or out_axes=(0, None)
Switch to jax.shard_map with out_specs=P(...) for name-based output sharding

Example fix

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

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

Defensive patterns

Strategy: type-guard

Validate before calling

from jax.tree_util import tree_leaves
assert all(isinstance(l, int) for l in tree_leaves(out_axes)), out_axes

Type guard

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

Prevention

out_axes mirrors in_axes: ints only
Use shard_map out_specs=P(...) for name-based output sharding

When it happens

Trigger: Calling smap with out_axes='i', out_axes=[0], or out_axes=(0, 'j').

Common situations: Symmetric confusion with in_axes: users specify mesh axis names or Python lists where integer output axes are expected; also copied-over PartitionSpec values from shard_map code.

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

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