jax-ml/jax · error · TypeError

shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.

Error message

shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.AbstractMesh` instance for its second argument, but got {mesh} of type {type(mesh)}.

What it means

shard_map's mesh argument (second argument) must be a jax.sharding.Mesh or jax.sharding.AbstractMesh instance. Passing anything else — dict, tuple, NamedSharding, a device list — raises this TypeError showing the offending value and its type.

Source

Thrown at jax/_src/shard_map.py:379

  return P(*[None] * axis + [axis_name])


def _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):
  if mesh is None:
    mesh = get_abstract_mesh()
    if mesh.empty:
      raise ValueError(
          "The context mesh cannot be empty. Use"
          " `jax.set_mesh(mesh)` to enter into a mesh context")
  else:
    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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create a real mesh: mesh = Mesh(jax.devices(), axis_names=('i',)) and pass that
2. If you already have a NamedSharding, extract its mesh or rebuild one over the same devices

Example fix

// before
jax.shard_map(f, mesh={'i': 4}, in_specs=P('i'), out_specs=P('i'))(x)

// after
mesh = jax.sharding.Mesh(jax.devices(), axis_names=('i',))
jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

from jax.sharding import Mesh, AbstractMesh
assert isinstance(mesh, (Mesh, AbstractMesh)), type(mesh)

def is_valid_mesh(m) -> bool:
    return isinstance(m, (jax.sharding.Mesh, jax.sharding.AbstractMesh))

When it happens

Trigger: Calling shard_map(f, mesh=P('i'), ...), mesh={'i': 4}, mesh=(devices,), or passing a jax.sharding.NamedSharding where the Mesh itself is expected.

Common situations: Confusing Mesh with NamedSharding (the latter is built FROM a mesh); constructing pseudo-meshes from raw device arrays; typos passing in_specs as the mesh argument positionally.

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

• The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t
• The context mesh {ctx_mesh} should match the mesh passed to
• shard_map requires a non-empty mesh. Got {mesh}
• jax.shard_map requires axis_names={axis_names} to be a subse
• Mapped away dimension of inputs passed to vmap should be sha

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