{"record":{"id":"773b1adeb2d145a1","repo":"jax-ml/jax","slug":"please-pass-jax-array-s-with-a-namedsharding-as","errorCode":null,"errorMessage":"Please pass `jax.Array`s with a `NamedSharding` as input to `shard_map` when passing `AbstractMesh` to the mesh argument.","messagePattern":"Please pass `jax\\.Array`s with a `NamedSharding` as input to `shard_map` when passing `AbstractMesh` to the mesh argument\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1200,"sourceCode":"    return str(map(names.get, range(aval.ndim)))\n  return ''\n\n# Eager evaluation\n\ndef get_mesh_from_args(args_flat, mesh):\n  for a in args_flat:\n    if (hasattr(a, 'sharding') and isinstance(a.sharding, NamedSharding)\n        and not a.sharding.mesh.is_scalar):  # pyrefly: ignore[missing-attribute]\n\n      if a.sharding.mesh.shape_tuple != mesh.shape_tuple:\n        aval = core.shaped_abstractify(a)\n        raise ValueError(\n            f\"Mesh shape of the input {a.sharding.mesh.shape_tuple} does not\"\n            \" match the mesh shape passed to shard_map \"\n            f\" {mesh.shape_tuple} for shape {aval.str_short()}\")\n      mesh = a.sharding.mesh\n  if isinstance(mesh, AbstractMesh):\n    raise ValueError(\n        \"Please pass `jax.Array`s with a `NamedSharding` as input to\"\n        \" `shard_map` when passing `AbstractMesh` to the mesh argument.\")\n  assert isinstance(mesh, Mesh)\n  return mesh\n\ndef _spec_to_vma(spec):\n  return frozenset(p for s in spec.partitions if s is not None\n                   for p in (s if isinstance(s, tuple) else (s,)))\n\ndef _mat_to_spec(mesh, mat):\n  return P(order_wrt_mesh(mesh, mat.varying), unreduced=mat.unreduced,\n           reduced=mat.reduced, unreduced_kind=mat.unreduced_kind)\n\ndef _spec_to_mat(spec) -> core.ManualAxisType:\n  return core.ManualAxisType(varying=_spec_to_vma(spec),\n                             unreduced=spec.unreduced, reduced=spec.reduced,\n                             unreduced_kind=spec.unreduced_kind)\n","sourceCodeStart":1182,"sourceCodeEnd":1218,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1182-L1218","documentation":"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.","triggerScenarios":"Calling shard_map(f, abstract_mesh, ...) where inputs are plain uncommitted numpy/jnp arrays without NamedSharding, or no array carries a concrete mesh.","commonSituations":"Using the new abstract-mesh API with raw host arrays; passing Python scalars/None placeholders; forgetting jax.device_put under the mesh context manager.","solutions":["device_put inputs with a NamedSharding over a concrete mesh before calling shard_map","Run under `with mesh:` and use jax.device_put, or construct arrays via operations that produce committed arrays","Pass a concrete Mesh instead of the AbstractMesh"],"exampleFix":"// before\ny = shard_map(f, amesh, np.arange(16.).reshape(4,4), in_specs=P('i'))\n// after\nx = jax.device_put(np.arange(16.).reshape(4,4), NamedSharding(backend_mesh, P('i')))\ny = shard_map(f, amesh, x, in_specs=P('i'))","handlingStrategy":"validation","validationCode":"def inputs_committed(args):\n    return all(hasattr(a, 'sharding') and isinstance(a.sharding, NamedSharding) for a in jax.tree.leaves(args))\n# then: jax.device_put(x, NamedSharding(concrete_mesh, P('i')))","typeGuard":"def all_named_sharded(args) -> bool:\n    return all(isinstance(getattr(a, 'sharding', None), NamedSharding) for a in jax.tree.leaves(args))","tryCatchPattern":null,"preventionTips":["Always device_put inputs under the mesh context before abstract-mesh shard_map","Prefer concrete Mesh when inputs are raw host arrays"],"tags":["shard-map","abstract-mesh","named-sharding","jax"],"backgroundTag":"missing-concrete-mesh","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}