{"record":{"id":"e04afb2e9440b661","repo":"jax-ml/jax","slug":"in-specs-passed-to-shard-map-s-does-not-match-t","errorCode":null,"errorMessage":"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.","messagePattern":"in_specs passed to shard_map: (.+?) does not match the specs of the input: (.+?) for arg: (.+?)\\. `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\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":301,"sourceCode":"    dyn_args      = [x for x, dyn in zip(args_flat, which_dyn) if dyn]\n    in_specs_flat = tuple(s for s, dyn in zip(in_specs_flat, which_dyn) if dyn)\n    dyn_argnums   = [i for i, dyn in enumerate(which_dyn) if dyn]\n    _check_specs_vs_args(f, mesh, in_tree, in_specs, dyn_argnums,\n                         in_specs_flat, dyn_args)\n\n    # TODO(yashkatariya): Add support for partial manual\n    mesh_axis_names_wo_vmap = (\n        frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names)\n    if (mesh_axis_names_wo_vmap == axis_names and\n        all(mesh._name_to_type[a] == AxisType.Explicit for a in axis_names)):\n      for a, s in zip(dyn_args, in_specs_flat):\n        if not isinstance(s, P): continue\n        arg_aval = typeof(a)\n        s = s._normalized_spec_for_aval(arg_aval.ndim)\n        if config.remove_size_one_mesh_axis_from_type.value:\n          s = remove_size_one_mesh_axis_from_spec(s, mesh)\n        if arg_aval.sharding.spec != s:\n          raise ValueError(\n              f\"in_specs passed to shard_map: {s} does not match the specs of\"\n              f\" the input: {arg_aval.sharding.spec} for arg: {typeof(a)}.\"\n              \" `in_specs` is an optional argument so you can omit specifying\"\n              \" it and shard_map will infer the in_specs from the arguments.\"\n              \" If you want to reshard your inputs, you can use `jax.reshard`\"\n              \" on the arguments and then pass those args to shard_map.\")\n\n    if (dbg.arg_names is not None and len(dyn_args) != len(dbg.arg_names)):\n      dbg = dbg.with_unknown_names()\n\n    def f_wrapped(*dyn_args):\n      dyn_args_iter = iter(dyn_args)\n      static_args_iter = iter(static_args)\n      all_args = [next(dyn_args_iter) if dyn else next(static_args_iter)\n                  for dyn in which_dyn]\n      args = tree_unflatten(in_tree, all_args)\n      ans = f(*args)\n      ans_ft = ft.flatten(ans)","sourceCodeStart":283,"sourceCodeEnd":319,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L283-L319","documentation":"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).","triggerScenarios":"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).","commonSituations":"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).","solutions":["Omit in_specs entirely and let shard_map infer specs from the argument sharding","Pre-shard inputs with jax.reshard(x, jax.sharding.NamedSharding(mesh, P('i'))) before the call","Create the arrays directly with the desired NamedSharding (e.g. jax.device_put with the target sharding)"],"exampleFix":"// before\nx = jnp.ones((8, 8))\nout = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)\n\n// after\nx = jax.reshard(x, jax.sharding.NamedSharding(mesh, P('i')))\nout = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)","handlingStrategy":"validation","validationCode":"def ensure_sharding(x, mesh, spec):\n    target = jax.sharding.NamedSharding(mesh, spec)\n    if not isinstance(getattr(x, 'sharding', None), type(target)) or x.sharding != target:\n        x = jax.reshard(x, target)\n    return x\n\nx = ensure_sharding(x, mesh, P('i'))\nout = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)","typeGuard":null,"tryCatchPattern":"try:\n    out = shmapped(x)\nexcept ValueError as e:\n    if 'does not match the specs of the input' in str(e):\n        x = jax.reshard(x, jax.sharding.NamedSharding(mesh, P('i')))\n        out = shmapped(x)\n    else: raise","preventionTips":["Create inputs directly with the target NamedSharding via device_put","Or omit in_specs and let shard_map infer","Wrap free-form inputs with jax.reshard at API boundaries"],"tags":["jax","shard-map","sharding","partition-spec","distributed"],"backgroundTag":"sharding-layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}