jax-ml/jax · critical · IndivisibleError

{self} implies that array axis {dim} is partitioned {p} time

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

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.

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)

Solutions

1. Pad the array so every sharded dim is divisible by its mesh axis size (e.g. pad batch to multiple of number of devices)
2. Adjust the PartitionSpec to shard a dimension that is divisible, or use P(None) for awkward dims
3. 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')))

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

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

• Mapped away dimension of inputs passed to vmap should be sha
• The 'sharding' attribute is not available on {self._error_re
• The is_fully_addressable property was called on {self._error
• The error code state and the predicate must be on the same m
• AbstractMesh should be of the same size across all shardings

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