jax-ml/jax · error · TypeError

`axis_names` argument of shard_map should be of type `frozen

Error message

`axis_names` argument of shard_map should be of type `frozenset` or `set`. Got type: {type(axis_names)}

What it means

shard_map's internal _shmap_checks requires axis_names to be a frozenset or set of mesh axis names. Passing any other container type — list, tuple, string, or None — raises this TypeError reporting the actual type received.

Source

Thrown at jax/_src/shard_map.py:390

    ctx_mesh = get_abstract_mesh()
    if not ctx_mesh.empty and mesh.abstract_mesh != ctx_mesh:
      raise ValueError(
          f"The context mesh {ctx_mesh} should match the mesh passed to"
          f" shard_map {mesh}")

  if not isinstance(mesh, (Mesh, AbstractMesh)):
    raise TypeError("shard_map requires a `jax.sharding.Mesh` or a "
                    "`jax.sharding.AbstractMesh` instance for its "
                    f"second argument, but got {mesh} of type {type(mesh)}.")
  if mesh.empty:
    raise ValueError(f"shard_map requires a non-empty mesh. Got {mesh}")

  mesh_axis_names_wo_vmap = (
      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names
  )

  if not isinstance(axis_names, (frozenset, set)):
    raise TypeError(
        "`axis_names` argument of shard_map should be of type `frozenset` or"
        f" `set`. Got type: {type(axis_names)}")
  if isinstance(axis_names, set):
    axis_names = frozenset(axis_names)
  if not axis_names:
    axis_names = mesh_axis_names_wo_vmap
  if not axis_names.issubset(mesh_axis_names_wo_vmap):
    raise ValueError(
        f"jax.shard_map requires axis_names={axis_names} to be a subset of "
        f"mesh.axis_names={mesh_axis_names_wo_vmap}")

  if (in_specs is Infer and
      not all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):
    axis_types = ', '.join(str(mesh._name_to_type[a]) for a in axis_names)
    if _smap:
      msg = (f"in_axes was not specified when axis_name={axis_names} was of"
             f" type {axis_types}")
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap names in a set/frozenset: axis_names=frozenset({'i','j'})
2. Use the public jax.shard_map API which normalizes the argument for you

Example fix

// before
shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'), axis_names=('i',))

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

assert isinstance(axis_names, (frozenset, set)), type(axis_names)
axis_names = frozenset(axis_names)

def is_valid_axis_names(ns) -> bool:
    return isinstance(ns, (frozenset, set))

When it happens

Trigger: Calling shard_map(..., axis_names=('i','j')) (tuple), axis_names=['i'] (list), axis_names='i' (bare string), or axis_names=None. This typically happens when calling the private _shard_map path or smap incorrectly, since the public shard_map API wraps names into a frozenset.

Common situations: Using jax.experimental.smap or internal/lower-level shard_map entry points; refactoring code that previously passed a single name; version upgrades that tightened axis_names typing.

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

• shard_map requires a callable for its first argument, but go
• shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.
• jax.shard_map requires axis_names={axis_names} to be a subse
• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li

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