{"record":{"id":"5b6a271090118d08","repo":"jax-ml/jax","slug":"sharding-spec-spec-implies-that-array-axis-dim","errorCode":null,"errorMessage":"Sharding spec {spec} implies that array axis {dim} is partitioned {size} times, but does not evenly divide the dimension size {sh}. Got shape: {shape} and sharding {sharding}","messagePattern":"Sharding spec (.+?) implies that array axis (.+?) is partitioned (.+?) times, but does not evenly divide the dimension size (.+?)\\. Got shape: (.+?) and sharding (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/core.py","lineNumber":2273,"sourceCode":"  else:\n    out = sharding.update(spec=modify_spec_for_auto_manual(\n        sharding.spec, sharding.mesh))\n  if config.remove_size_one_mesh_axis_from_type.value:\n    out = out.update(spec=ns.remove_size_one_mesh_axis_from_spec(out.spec, out.mesh))\n  if len(out.spec) != ndim:\n    out = _make_lengths_same(out, ndim)\n  return out\n\ndef _check_divisibility(sharding, shape):\n  mesh = sharding.mesh\n  for dim, (spec, sh) in enumerate(zip(sharding.spec.partitions, shape)):\n    if spec is None:\n      continue\n    spec = spec if isinstance(spec, tuple) else (spec,)\n    size = math.prod(mesh.shape[s] for s in spec)\n    _, remainder = divmod(sh, size)\n    if remainder != 0:\n      raise ValueError(\n          f\"Sharding spec {spec} implies that array axis {dim} is partitioned\"\n          f\" {size} times, but does not evenly divide the dimension size {sh}.\"\n          f\" Got shape: {shape} and sharding {sharding}\")\n\n@cache(max_size=4096,\n       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)\ndef get_sharding(sharding, shape):\n  \"\"\"Modifies and checks the sharding.\n\n  Some modifications/checks include:\n    * Making the length of specs the same as ndim\n    * If a mesh axis is mentioned in pspec is Auto/Manual, replace it with None\n    * Checking for len(spec)-ndim match\n    * Checking if the mesh is an AbstractMesh.\n  \"\"\"\n  ndim = len(shape)\n  if sharding is None:\n    return _empty_sharding(ndim)","sourceCodeStart":2255,"sourceCodeEnd":2291,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2255-L2291","documentation":"Raised when building a sharded array type: an axis of the array is partitioned across a mesh dimension product that does not evenly divide that axis's size, producing an impossible/remainder sharding.","triggerScenarios":"Using NamedSharding/PartitionSpec (or GSPMD sharding annotations) where the product of mesh dimension sizes assigned to an array axis does not divide that axis's length, e.g., sharding a size-3 axis over a 2-device mesh dimension.","commonSituations":"Mismatch between array shape and jax.sharding.Mesh shape; leftover partitions from an older mesh; using pmap with more devices than a dimension's size; sharding tiny axes (size 1 or primes) across multi-device mesh axes.","solutions":["Adjust the PartitionSpec so the product of mesh dims for each axis divides that axis's size (or use None to replicate)","Pad or reshape the array so the sharded dimension is divisible by the mesh sub-shape","Resize/reorder the mesh (mesh_shape) so each partitioned axis divides evenly","Compute divisibility programmatically before constructing the sharding (see validation code)"],"exampleFix":"# before\nmesh = Mesh(jax.devices(), ('x',))\nsharding = NamedSharding(mesh, P('x',))  # shape (3, 4) over 8 devices -> error\n\n# after\nmesh = Mesh(jax.devices()[:4], ('x',))  # axis 0 size 4, wait 3 not divisible either\n# correct: shard axis of size 4 over 4 devices\nsharding = NamedSharding(mesh, P(None, 'x'))  # shape (3, 4), shard axis 1 of size 4 over 4","handlingStrategy":"validation","validationCode":"import math\ndef sharding_divides(shape, mesh_shape, spec):\n    for dim, s in enumerate(spec):\n        if s is None: continue\n        s = (s,) if not isinstance(s, tuple) else s\n        size = math.prod(mesh_shape[axis] for axis in s)\n        if shape[dim] % size != 0:\n            return False\n    return True\nassert sharding_divides(x.shape, mesh.shape, P('data',).specs)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive PartitionSpec programmatically from mesh shape and array shape","Assert divisibility of each sharded axis before constructing NamedSharding","Pad prime-sized or odd axes or mark them replicated (None) in the spec"],"tags":["jax","sharding","named-sharding","mesh","partitionspec"],"backgroundTag":"jax-uneven-sharding","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}