{"record":{"id":"5cd7089a2fcbcd11","repo":"jax-ml/jax","slug":"self-implies-that-array-axis-dim-is-partitione","errorCode":null,"errorMessage":"{self} implies that array axis {dim} is partitioned {p} times, but the dimension size is {s} (full shape: {global_shape}, per-dimension tiling factors: {tuple(partitions)} should evenly divide the shape)","messagePattern":"(.+?) implies that array axis (.+?) is partitioned (.+?) times, but the dimension size is (.+?) \\(full shape: (.+?), per-dimension tiling factors: (.+?) should evenly divide the shape\\)","errorType":"validation","errorClass":"IndivisibleError","httpStatus":null,"severity":"critical","filePath":"jax/_src/sharding.py","lineNumber":72,"sourceCode":"  indices = op_sharding_to_indices(hlo_sharding, global_shape,\n                                   len(s._device_assignment))\n  return dict(safe_zip(s._device_assignment, indices))\n\n\n@cache(max_size=4096, trace_context_in_key=False)\ndef _common_shard_shape(self, global_shape: Shape) -> Shape:\n  hlo_sharding = self._to_xla_hlo_sharding(len(global_shape))\n  if is_hlo_sharding_replicated(hlo_sharding):\n    return global_shape\n  if hlo_sharding.is_unreduced():\n    return global_shape\n  partitions, _ = get_num_ways_dim_sharded(hlo_sharding)\n  assert len(partitions) == len(global_shape), (len(partitions), len(global_shape))\n  out = []\n  for dim, (s, p) in enumerate(safe_zip(global_shape, partitions)):\n    quotient, remainder = divmod(s, p)\n    if remainder != 0:\n      raise IndivisibleError(\n          f\"{self} implies that array axis {dim} is partitioned \"\n          f\"{p} times, but the dimension size is {s} \"\n          f\"(full shape: {global_shape}, \"\n          f\"per-dimension tiling factors: {tuple(partitions)} should evenly \"\n          \"divide the shape)\")\n    out.append(quotient)\n  return tuple(out)\n\ndef common_is_equivalent_to(s1: Sharding, s2: Sharding, ndim: int,\n                            check_devices: bool = True) -> bool:\n  hlo_s_eq = are_hlo_shardings_equal(\n      s1._to_xla_hlo_sharding(ndim), s2._to_xla_hlo_sharding(ndim))\n  mem_eq = s1.memory_kind == s2.memory_kind\n  if check_devices:\n    return (hlo_s_eq and mem_eq and\n            s1._internal_device_list == s2._internal_device_list)\n  else:\n    return hlo_s_eq and mem_eq","sourceCodeStart":54,"sourceCodeEnd":90,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/sharding.py#L54-L90","documentation":"Raised as IndivisibleError when computing per-shard local shapes: the HloSharding shards some array dimension p times but that dimension's size is not divisible by p, so even tiling is impossible.","triggerScenarios":"Sharding a shape-(3,) array with a mesh whose 'data' axis has 4 devices (P('data')), or NamedSharding address_indices/shard_shape on such a layout; per-dimension tiling factors from get_num_ways_dim_sharded don't divide the shape.","commonSituations":"Shape not divisible by device count (e.g. batch 1000 across 8 GPUs is fine, 1001 fails); last-batch sharding in distributed training; padding lost during preprocessing.","solutions":["Pad the array so every sharded dim is divisible by its mesh axis size (e.g. pad batch to multiple of number of devices)","Adjust the PartitionSpec to shard a dimension that is divisible, or use P(None) for awkward dims","Change mesh shape / number of devices so the divisor fits"],"exampleFix":"# before\nsharded = device_put(x, NamedSharding(mesh, P('data')))  # x.shape[0]=1001, axis size 8\n\n# after\npad = (-x.shape[0]) % 8\nx = jnp.pad(x, ((0, pad),) + ((0,0),)*(x.ndim-1))\nsharded = device_put(x, NamedSharding(mesh, P('data')))","handlingStrategy":"validation","validationCode":"import math\nfor dim, size in enumerate(x.shape):\n    ways = prod(mesh.shape[n] for n in (spec[dim] if isinstance(spec[dim], tuple) else ((spec[dim],) if spec[dim] else ())) )\n    assert size % ways == 0 if ways else True, f'dim {dim} size {size} not divisible by {ways}'","typeGuard":null,"tryCatchPattern":"try:\n    y = jax.device_put(x, sharding)\nexcept jax._src.sharding_impls.IndivisibleError:\n    pad = [-s % w for s, w in zip(x.shape, ways)]\n    y = jax.device_put(jnp.pad(x, ...), sharding)","preventionTips":["Pad batches to a multiple of the data-parallel device count in your input pipeline","Assert divisibility shape vs mesh axis sizes in tests","Prefer sharding dims that are statically divisible (hidden dims, vocab)"],"tags":["jax","sharding","indivisible-shape","pjit","distributed"],"backgroundTag":"array-shape-not-divisible-by-shards","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}