jax-ml/jax · error · ValueError

Please pass `jax.Array`s with a `NamedSharding` as input to

Error message

Please pass `jax.Array`s with a `NamedSharding` as input to `shard_map` when passing `AbstractMesh` to the mesh argument.

What it means

When an AbstractMesh is passed as the mesh argument, shard_map must concretize it from the arguments; it requires all (relevant) inputs to be jax.Arrays with NamedShardings so a concrete Mesh can be recovered. Otherwise it raises this ValueError.

Source

Thrown at jax/_src/shard_map.py:1200

    return str(map(names.get, range(aval.ndim)))
  return ''

# Eager evaluation

def get_mesh_from_args(args_flat, mesh):
  for a in args_flat:
    if (hasattr(a, 'sharding') and isinstance(a.sharding, NamedSharding)
        and not a.sharding.mesh.is_scalar):  # pyrefly: ignore[missing-attribute]

      if a.sharding.mesh.shape_tuple != mesh.shape_tuple:
        aval = core.shaped_abstractify(a)
        raise ValueError(
            f"Mesh shape of the input {a.sharding.mesh.shape_tuple} does not"
            " match the mesh shape passed to shard_map "
            f" {mesh.shape_tuple} for shape {aval.str_short()}")
      mesh = a.sharding.mesh
  if isinstance(mesh, AbstractMesh):
    raise ValueError(
        "Please pass `jax.Array`s with a `NamedSharding` as input to"
        " `shard_map` when passing `AbstractMesh` to the mesh argument.")
  assert isinstance(mesh, Mesh)
  return mesh

def _spec_to_vma(spec):
  return frozenset(p for s in spec.partitions if s is not None
                   for p in (s if isinstance(s, tuple) else (s,)))

def _mat_to_spec(mesh, mat):
  return P(order_wrt_mesh(mesh, mat.varying), unreduced=mat.unreduced,
           reduced=mat.reduced, unreduced_kind=mat.unreduced_kind)

def _spec_to_mat(spec) -> core.ManualAxisType:
  return core.ManualAxisType(varying=_spec_to_vma(spec),
                             unreduced=spec.unreduced, reduced=spec.reduced,
                             unreduced_kind=spec.unreduced_kind)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. device_put inputs with a NamedSharding over a concrete mesh before calling shard_map
2. Run under `with mesh:` and use jax.device_put, or construct arrays via operations that produce committed arrays
3. Pass a concrete Mesh instead of the AbstractMesh

Example fix

// before
y = shard_map(f, amesh, np.arange(16.).reshape(4,4), in_specs=P('i'))
// after
x = jax.device_put(np.arange(16.).reshape(4,4), NamedSharding(backend_mesh, P('i')))
y = shard_map(f, amesh, x, in_specs=P('i'))

def inputs_committed(args):
    return all(hasattr(a, 'sharding') and isinstance(a.sharding, NamedSharding) for a in jax.tree.leaves(args))
# then: jax.device_put(x, NamedSharding(concrete_mesh, P('i')))

def all_named_sharded(args) -> bool:
    return all(isinstance(getattr(a, 'sharding', None), NamedSharding) for a in jax.tree.leaves(args))

When it happens

Trigger: Calling shard_map(f, abstract_mesh, ...) where inputs are plain uncommitted numpy/jnp arrays without NamedSharding, or no array carries a concrete mesh.

Common situations: Using the new abstract-mesh API with raw host arrays; passing Python scalars/None placeholders; forgetting jax.device_put under the mesh context manager.

Related errors

• Mesh shape of the input {a.sharding.mesh.shape_tuple} does n
• Mesh must be provided for shard_map with checkify.
• Unsupported aval type: {type(v)}
• Sharding spec {spec} implies that array axis {dim} is partit
• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of

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