jax-ml/jax · error · ValueError

shard_map in_specs divisibility error (msg from _spec_divisi

Error message

shard_map in_specs divisibility error (msg from _spec_divisibility_error)

What it means

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.

Source

Thrown at jax/_src/shard_map.py:541

def _check_specs_vs_args(
    f: Callable, mesh: Mesh | AbstractMesh, in_tree: PyTreeDef, in_specs: Specs,
    dyn_argnums: Sequence[int], in_specs_flat: Sequence[P],
    xs: Sequence) -> None:
  in_avals = map(core.shaped_abstractify, xs)
  fail = [a if isinstance(p, P) and len(p) > a.ndim else no_fail
          for p, a in zip(in_specs_flat, in_avals)]
  if any(f is not no_fail for f in fail):
    fail = _expand_fail(in_tree, dyn_argnums, fail)
    msg = _spec_rank_error(SpecErrorType.input, f, in_tree, in_specs, fail)
    raise ValueError(msg)
  bad = lambda a, d, ns: a.shape[d] % prod(mesh.shape[n] for n in ns)
  fail = [a if (isinstance(s, P) and
                any(bad(a, d, ns) for d, ns in _spec_to_names(s).items()))
          else no_fail for a, s in zip(in_avals, in_specs_flat)]
  if any(f is not no_fail for f in fail):
    fail = _expand_fail(in_tree, dyn_argnums, fail)
    msg = _spec_divisibility_error(f, mesh, in_tree, in_specs, fail)
    raise ValueError(msg)

def _expand_fail(in_tree: PyTreeDef, dyn_argnums: Sequence[int],
                 fail: Sequence[core.ShapedArray | NoFail]
                 ) -> list[core.ShapedArray | NoFail]:
  fail_: list[core.ShapedArray | NoFail] = [no_fail] * in_tree.num_leaves
  for i, f in zip(dyn_argnums, fail):
    fail_[i] = f
  return fail_

def _spec_rank_error(
    error_type: SpecErrorType, f: Callable, tree: PyTreeDef, specs: Specs,
    fails: list[core.ShapedArray | NoFail]) -> str:
  fun_name = util_fun_name(f)
  if error_type == SpecErrorType.input:
    prefix, base = 'in', 'the passed args'
    ba = _try_infer_args(f, tree)
  else:
    prefix, base = 'out', f'{fun_name}(*args)'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad or resize the array dimension to a multiple of the product of mapped axis sizes
2. Change the spec so the dimension is replicated (None) or mapped to fewer/smaller axes
3. Choose mesh axis sizes that divide the corresponding array dimensions

Example fix

// before
shard_map(f, mesh, jnp.ones(10), in_specs=P('x'))  # mesh x=4
// after
shard_map(f, mesh, jnp.ones(12), in_specs=P('x'))

from math import prod
def divisible(tree, args, specs, mesh):
    for a, p in zip(map(jax.core.shaped_abstractify, jax.tree.leaves(args)), jax.tree.leaves(specs)):
        for d, ns in _spec_to_names(p).items():
            assert a.shape[d] % prod(mesh.shape[n] for n in ns) == 0, (a.shape, ns)

try: shard_map(...) except ValueError as e: if 'divisibility' in str(e).lower() or 'not divisible' in str(e): pad inputs; else: raise

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• Block size must be a multiple of the input size. Got block {
• Every block dimension must be either a multiple or factor of
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac

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