jax-ml/jax · error · ValueError

in_specs passed to shard_map: {s} does not match the specs o

Error message

in_specs passed to shard_map: {s} does not match the specs of the input: {arg_aval.sharding.spec} for arg: {typeof(a)}. `in_specs` is an optional argument so you can omit specifying it and shard_map will infer the in_specs from the arguments. If you want to reshard your inputs, you can use `jax.reshard` on the arguments and then pass those args to shard_map.

What it means

When shard_map receives explicit in_specs, JAX checks that each input array's actual sharding (its Sharding spec) matches the requested PartitionSpec. A mismatch raises this ValueError, which also suggests two remedies: omit in_specs to infer from the arguments, or pre-shard with jax.reshard. This check only fires when the mismatch is not normalized away (e.g. size-1 mesh axis removal is config-dependent).

Source

Thrown at jax/_src/shard_map.py:301

    dyn_args      = [x for x, dyn in zip(args_flat, which_dyn) if dyn]
    in_specs_flat = tuple(s for s, dyn in zip(in_specs_flat, which_dyn) if dyn)
    dyn_argnums   = [i for i, dyn in enumerate(which_dyn) if dyn]
    _check_specs_vs_args(f, mesh, in_tree, in_specs, dyn_argnums,
                         in_specs_flat, dyn_args)

    # TODO(yashkatariya): Add support for partial manual
    mesh_axis_names_wo_vmap = (
        frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names)
    if (mesh_axis_names_wo_vmap == axis_names and
        all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):
      for a, s in zip(dyn_args, in_specs_flat):
        if not isinstance(s, P): continue
        arg_aval = typeof(a)
        s = s._normalized_spec_for_aval(arg_aval.ndim)
        if config.remove_size_one_mesh_axis_from_type.value:
          s = remove_size_one_mesh_axis_from_spec(s, mesh)
        if arg_aval.sharding.spec != s:
          raise ValueError(
              f"in_specs passed to shard_map: {s} does not match the specs of"
              f" the input: {arg_aval.sharding.spec} for arg: {typeof(a)}."
              " `in_specs` is an optional argument so you can omit specifying"
              " it and shard_map will infer the in_specs from the arguments."
              " If you want to reshard your inputs, you can use `jax.reshard`"
              " on the arguments and then pass those args to shard_map.")

    if (dbg.arg_names is not None and len(dyn_args) != len(dbg.arg_names)):
      dbg = dbg.with_unknown_names()

    def f_wrapped(*dyn_args):
      dyn_args_iter = iter(dyn_args)
      static_args_iter = iter(static_args)
      all_args = [next(dyn_args_iter) if dyn else next(static_args_iter)
                  for dyn in which_dyn]
      args = tree_unflatten(in_tree, all_args)
      ans = f(*args)
      ans_ft = ft.flatten(ans)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Omit in_specs entirely and let shard_map infer specs from the argument sharding
2. Pre-shard inputs with jax.reshard(x, jax.sharding.NamedSharding(mesh, P('i'))) before the call
3. Create the arrays directly with the desired NamedSharding (e.g. jax.device_put with the target sharding)

Example fix

// before
x = jnp.ones((8, 8))
out = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

// after
x = jax.reshard(x, jax.sharding.NamedSharding(mesh, P('i')))
out = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

def ensure_sharding(x, mesh, spec):
    target = jax.sharding.NamedSharding(mesh, spec)
    if not isinstance(getattr(x, 'sharding', None), type(target)) or x.sharding != target:
        x = jax.reshard(x, target)
    return x

x = ensure_sharding(x, mesh, P('i'))
out = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)

try:
    out = shmapped(x)
except ValueError as e:
    if 'does not match the specs of the input' in str(e):
        x = jax.reshard(x, jax.sharding.NamedSharding(mesh, P('i')))
        out = shmapped(x)
    else: raise

When it happens

Trigger: Calling shard_map with in_specs=P('i') on an array that is actually replicated or sharded differently (e.g. created on a single device while mesh has multiple devices, or sharded along another axis).

Common situations: Passing locally-created jnp.ones(...) arrays (single-device) into a multi-device shard_map; mixing NamedSharding layouts; upgrading JAX versions where this strictness was introduced (previously mismatched inputs were implicitly resharded).

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Mesh must be provided for shard_map with checkify.
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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