jax-ml/jax · error · ValueError

Specified axis_size {axis_size} doesn't match received axis_

Error message

Specified axis_size {axis_size} doesn't match received axis_size {local_axis_size}.

What it means

In single-process mode, pmap validates that an explicitly supplied `axis_size` equals the mapped-axis size inferred from the arguments. A mismatch means the declared parallel width contradicts the actual batch dimension of the inputs.

Source

Thrown at jax/_src/pmap.py:691

  Args:
    devices: The mesh devices tuple.
    backend: The backend to use.
    local_axis_size: The local axis size (per-process).
    axis_size: User-specified global axis size (optional).
    trace_state_clean: True if in execution mode (not tracing).

  Returns:
    Tuple of effective mesh devices sliced appropriately.

  Raises:
    ValueError: If axis_size doesn't match inferred size in single-process.
  """
  process_count = xb.process_count(backend)

  # Validate explicit axis_size in single-process mode
  if (process_count == 1 and axis_size is not None and
      axis_size != local_axis_size):
    raise ValueError(
        f"Specified axis_size {axis_size} doesn't match received "
        f"axis_size {local_axis_size}.")

  # Compute global_axis_size
  if axis_size is not None:
    global_axis_size = axis_size
  elif process_count > 1:
    global_axis_size = local_axis_size * process_count
    # Validate all processes have the same number of local devices
    assert all(
        len(xb.local_devices(pi, backend)) == xb.local_device_count(backend)
        for pi in range(process_count))
  else:
    global_axis_size = local_axis_size

  # Determine mesh devices
  if devices is not None:
    mesh_devices = devices

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Drop the explicit axis_size and let pmap infer it from arguments
2. Resize the input's mapped axis to match axis_size (pad or rebatch)
3. Compute axis_size dynamically from local device count and batch size

Example fix

# before
f = jax.pmap(fn, axis_size=8); f(x)  # x.shape[0] == 4
# after
f = jax.pmap(fn)
f(x)

local_size = x.shape[in_axis]
if axis_size is not None:
    assert axis_size == local_size, f'{axis_size=} != {local_size=}'

When it happens

Trigger: `jax.pmap(f, axis_size=8)(x)` where x's mapped axis (per in_axes) is not 8, e.g. a batch of size 4 on an 8-device local run.

Common situations: Hardcoding device counts that differ from the dataset batch size; slicing data after pmap was configured; multi-process code run in single-process test mode.

Related errors

• pmap requires at least one argument with a mapped axis.
• {name} wrapped function must be passed at least one argument
• {name} must have at least one non-None value in in_axes or a
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out

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