{"record":{"id":"34720f11fac2e8ba","repo":"jax-ml/jax","slug":"specified-axis-size-axis-size-doesn-t-match-rece","errorCode":null,"errorMessage":"Specified axis_size {axis_size} doesn't match received axis_size {local_axis_size}.","messagePattern":"Specified axis_size (.+?) doesn't match received axis_size (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pmap.py","lineNumber":691,"sourceCode":"  Args:\n    devices: The mesh devices tuple.\n    backend: The backend to use.\n    local_axis_size: The local axis size (per-process).\n    axis_size: User-specified global axis size (optional).\n    trace_state_clean: True if in execution mode (not tracing).\n\n  Returns:\n    Tuple of effective mesh devices sliced appropriately.\n\n  Raises:\n    ValueError: If axis_size doesn't match inferred size in single-process.\n  \"\"\"\n  process_count = xb.process_count(backend)\n\n  # Validate explicit axis_size in single-process mode\n  if (process_count == 1 and axis_size is not None and\n      axis_size != local_axis_size):\n    raise ValueError(\n        f\"Specified axis_size {axis_size} doesn't match received \"\n        f\"axis_size {local_axis_size}.\")\n\n  # Compute global_axis_size\n  if axis_size is not None:\n    global_axis_size = axis_size\n  elif process_count > 1:\n    global_axis_size = local_axis_size * process_count\n    # Validate all processes have the same number of local devices\n    assert all(\n        len(xb.local_devices(pi, backend)) == xb.local_device_count(backend)\n        for pi in range(process_count))\n  else:\n    global_axis_size = local_axis_size\n\n  # Determine mesh devices\n  if devices is not None:\n    mesh_devices = devices","sourceCodeStart":673,"sourceCodeEnd":709,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pmap.py#L673-L709","documentation":"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.","triggerScenarios":"`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.","commonSituations":"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.","solutions":["Drop the explicit axis_size and let pmap infer it from arguments","Resize the input's mapped axis to match axis_size (pad or rebatch)","Compute axis_size dynamically from local device count and batch size"],"exampleFix":"# before\nf = jax.pmap(fn, axis_size=8); f(x)  # x.shape[0] == 4\n# after\nf = jax.pmap(fn)\nf(x)","handlingStrategy":"validation","validationCode":"local_size = x.shape[in_axis]\nif axis_size is not None:\n    assert axis_size == local_size, f'{axis_size=} != {local_size=}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Omit axis_size and let pmap infer","Derive axis_size from data/device counts at runtime, not constants"],"tags":["jax","pmap","axis-size","shape-mismatch"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}