jax-ml/jax · critical · IndivisibleError

implies that array axis is partitioned times, but the…

Error message

{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)

What it means

Validation inside Sharding._common_shard_shape: the per-dimension tiling factors implied by the HLO sharding must evenly divide the corresponding global shape dimensions. It fires when an axis is partitioned p times but its dimension size s is not divisible by p, so a consistent per-shard shape cannot be computed (surfaced via shard_shape).

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

Example fix

# before
sharded = device_put(x, NamedSharding(mesh, P('data')))  # x.shape[0]=1001, axis size 8

# after
pad = (-x.shape[0]) % 8
x = jnp.pad(x, ((0, pad),) + ((0,0),)*(x.ndim-1))
sharded = device_put(x, NamedSharding(mesh, P('data')))

Defensive patterns

Strategy: validation

Validate before calling

import math
for dim, size in enumerate(x.shape):
    ways = prod(mesh.shape[n] for n in (spec[dim] if isinstance(spec[dim], tuple) else ((spec[dim],) if spec[dim] else ())) )
    assert size % ways == 0 if ways else True, f'dim {dim} size {size} not divisible by {ways}'

Try / catch

try:
    y = jax.device_put(x, sharding)
except jax._src.sharding_impls.IndivisibleError:
    pad = [-s % w for s, w in zip(x.shape, ways)]
    y = jax.device_put(jnp.pad(x, ...), sharding)

Prevention

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)

When it happens

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

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

Related errors

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

Appendix: source

Thrown at jax/_src/sharding.py:72

  indices = op_sharding_to_indices(hlo_sharding, global_shape,
                                   len(s._device_assignment))
  return dict(safe_zip(s._device_assignment, indices))


@cache(max_size=4096, trace_context_in_key=False)
def _common_shard_shape(self, global_shape: Shape) -> Shape:
  hlo_sharding = self._to_xla_hlo_sharding(len(global_shape))
  if is_hlo_sharding_replicated(hlo_sharding):
    return global_shape
  if hlo_sharding.is_unreduced():
    return global_shape
  partitions, _ = get_num_ways_dim_sharded(hlo_sharding)
  assert len(partitions) == len(global_shape), (len(partitions), len(global_shape))
  out = []
  for dim, (s, p) in enumerate(safe_zip(global_shape, partitions)):
    quotient, remainder = divmod(s, p)
    if remainder != 0:
      raise IndivisibleError(
          f"{self} implies that array axis {dim} is partitioned "
          f"{p} times, but the dimension size is {s} "
          f"(full shape: {global_shape}, "
          f"per-dimension tiling factors: {tuple(partitions)} should evenly "
          "divide the shape)")
    out.append(quotient)
  return tuple(out)

def common_is_equivalent_to(s1: Sharding, s2: Sharding, ndim: int,
                            check_devices: bool = True) -> bool:
  hlo_s_eq = are_hlo_shardings_equal(
      s1._to_xla_hlo_sharding(ndim), s2._to_xla_hlo_sharding(ndim))
  mem_eq = s1.memory_kind == s2.memory_kind
  if check_devices:
    return (hlo_s_eq and mem_eq and
            s1._internal_device_list == s2._internal_device_list)
  else:
    return hlo_s_eq and mem_eq

View on GitHub (pinned to 1e1c6a8fc0)