{"record":{"id":"b995e9405817d1e9","repo":"jax-ml/jax","slug":"shard-map-in-specs-divisibility-error-msg-from-s","errorCode":null,"errorMessage":"shard_map in_specs divisibility error (msg from _spec_divisibility_error)","messagePattern":"shard_map in_specs divisibility error \\(msg from _spec_divisibility_error\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":541,"sourceCode":"def _check_specs_vs_args(\n    f: Callable, mesh: Mesh | AbstractMesh, in_tree: PyTreeDef, in_specs: Specs,\n    dyn_argnums: Sequence[int], in_specs_flat: Sequence[P],\n    xs: Sequence) -> None:\n  in_avals = map(core.shaped_abstractify, xs)\n  fail = [a if isinstance(p, P) and len(p) > a.ndim else no_fail\n          for p, a in zip(in_specs_flat, in_avals)]\n  if any(f is not no_fail for f in fail):\n    fail = _expand_fail(in_tree, dyn_argnums, fail)\n    msg = _spec_rank_error(SpecErrorType.input, f, in_tree, in_specs, fail)\n    raise ValueError(msg)\n  bad = lambda a, d, ns: a.shape[d] % prod(mesh.shape[n] for n in ns)\n  fail = [a if (isinstance(s, P) and\n                any(bad(a, d, ns) for d, ns in _spec_to_names(s).items()))\n          else no_fail for a, s in zip(in_avals, in_specs_flat)]\n  if any(f is not no_fail for f in fail):\n    fail = _expand_fail(in_tree, dyn_argnums, fail)\n    msg = _spec_divisibility_error(f, mesh, in_tree, in_specs, fail)\n    raise ValueError(msg)\n\ndef _expand_fail(in_tree: PyTreeDef, dyn_argnums: Sequence[int],\n                 fail: Sequence[core.ShapedArray | NoFail]\n                 ) -> list[core.ShapedArray | NoFail]:\n  fail_: list[core.ShapedArray | NoFail] = [no_fail] * in_tree.num_leaves\n  for i, f in zip(dyn_argnums, fail):\n    fail_[i] = f\n  return fail_\n\ndef _spec_rank_error(\n    error_type: SpecErrorType, f: Callable, tree: PyTreeDef, specs: Specs,\n    fails: list[core.ShapedArray | NoFail]) -> str:\n  fun_name = util_fun_name(f)\n  if error_type == SpecErrorType.input:\n    prefix, base = 'in', 'the passed args'\n    ba = _try_infer_args(f, tree)\n  else:\n    prefix, base = 'out', f'{fun_name}(*args)'","sourceCodeStart":523,"sourceCodeEnd":559,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L523-L559","documentation":"Every array dimension referenced by a spec must be evenly divisible by the product of the mesh axis sizes mapping to it. Otherwise the per-device local shard size would be fractional; shard_map raises this ValueError listing the offending args and mesh.","triggerScenarios":"x of shape (10,) on mesh {'x': 4} with in_specs=P('x') since 10 % 4 != 0; or a dim of 12 mapped to two axes of size 3 and 4 is fine, but 10 to (4,) fails.","commonSituations":"Hard-coded array sizes not multiples of mesh sizes; changing device count (e.g. 8 GPUs -> 4) without resizing data; attention heads or hidden dims smaller than an axis.","solutions":["Pad or resize the array dimension to a multiple of the product of mapped axis sizes","Change the spec so the dimension is replicated (None) or mapped to fewer/smaller axes","Choose mesh axis sizes that divide the corresponding array dimensions"],"exampleFix":"// before\nshard_map(f, mesh, jnp.ones(10), in_specs=P('x'))  # mesh x=4\n// after\nshard_map(f, mesh, jnp.ones(12), in_specs=P('x'))","handlingStrategy":"validation","validationCode":"from math import prod\ndef divisible(tree, args, specs, mesh):\n    for a, p in zip(map(jax.core.shaped_abstractify, jax.tree.leaves(args)), jax.tree.leaves(specs)):\n        for d, ns in _spec_to_names(p).items():\n            assert a.shape[d] % prod(mesh.shape[n] for n in ns) == 0, (a.shape, ns)","typeGuard":null,"tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'divisibility' in str(e).lower() or 'not divisible' in str(e): pad inputs; else: raise","preventionTips":["Choose model dims as multiples of total mesh sizes per axis","Add a startup assert that array shapes divide mesh shapes"],"tags":["shard-map","divisibility","shape-mismatch","jax"],"backgroundTag":"shape-not-divisible-by-mesh","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}