{"record":{"id":"a1c01df332d86114","repo":"jax-ml/jax","slug":"name-was-requested-to-map-its-argument-along-axi","errorCode":null,"errorMessage":"{name} was requested to map its argument along axis {axis}, which implies that its rank should be at least {min_rank}, but is only {len(shape)} (its shape is {shape})","messagePattern":"(.+?) was requested to map its argument along axis (.+?), which implies that its rank should be at least (.+?), but is only (.+?) \\(its shape is (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1324,"sourceCode":"\n  def _get_axis_size(name: str, x, axis: int) -> core.AxisSize | None:\n    shape: tuple[core.AxisSize, ...] = ()\n    try:\n      shape = np.shape(x)\n      return shape[axis]\n    except (IndexError, TypeError) as e:\n      if not core.valid_jaxtype(x) or not isinstance(axis, int):\n        return None  # Suppress the check for custom vmappable types.\n      if core.typeof(x).is_high:\n        raise ValueError(\n            f\"{name} was requested to map a value of non-array type \"\n            f\"{core.typeof(x)} along axis {axis}, but non-array types can't \"\n            \"be mapped along an integer axis. Instead pass a mapping spec (a \"\n            \"MappingSpec instance) as this argument's in_axes entry, and \"\n            \"pass axis_size explicitly.\") from None\n      min_rank = axis + 1 if axis >= 0 else -axis\n      # TODO(mattjj): better error message here\n      raise ValueError(\n          f\"{name} was requested to map its argument along axis {axis}, \"\n          f\"which implies that its rank should be at least {min_rank}, \"\n          f\"but is only {len(shape)} (its shape is {shape})\") from e\n\n  all_mapped_sizes = [\n    None if d is None else _get_axis_size(name, x, d)\n    for x, d in zip(vals, dims)\n  ]\n  all_sizes = [s for s in all_mapped_sizes if s is not None]\n  if axis_size is not None:\n    all_sizes.append(axis_size)\n  sizes = core.dedup_referents(all_sizes)\n  if len(sizes) == 1:\n    sz, = sizes\n    return sz\n  if not sizes:\n    raise ValueError(f\"{name} must have at least one non-None value in in_axes \"\n                     \"or axis_size must be specified\")","sourceCodeStart":1306,"sourceCodeEnd":1342,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1306-L1342","documentation":"Raised when vmap requests an in_axes index larger than the argument's rank: e.g. in_axes=2 for a 1-D array. The mapped axis must exist, so the array needs at least |axis|+1 dimensions.","triggerScenarios":"jax.vmap(f, in_axes=1)(jnp.zeros(4)) (mapping axis 1 of a rank-1 array); using negative in_axes like -2 on a 1-D input; wrong batch dimension index after refactoring shapes.","commonSituations":"Data pipeline changed an array's rank (squeeze/reshape) while in_axes stayed the same; off-by-one in the batch dimension index; nested vmaps with cumulative axis counts exceeding rank.","solutions":["Fix in_axes to a valid axis for the argument's actual rank","Reshape/expand the input so it has the expected rank (e.g. add a batch dim)","For nested vmaps remember inner in_axes count from the pre-outer-vmap perspective"],"exampleFix":"// before\njax.vmap(f, in_axes=1)(jnp.zeros(4))\n// after\njax.vmap(f, in_axes=0)(jnp.zeros((4, 4)))\n# or\njax.vmap(f, in_axes=1)(jnp.zeros(4)[:, None])","handlingStrategy":"validation","validationCode":"for leaf, ax in zip(tree_leaves(args), tree_leaves(in_axes)):\n    if isinstance(ax, int) and hasattr(leaf, 'ndim'):\n        assert -leaf.ndim <= ax < leaf.ndim, f'in_axes {ax} invalid for shape {leaf.shape}'","typeGuard":"def axis_valid(x, ax): return hasattr(x, 'ndim') and -x.ndim <= ax < x.ndim","tryCatchPattern":null,"preventionTips":["Assert in_axes within each argument's rank before vmap","Log shapes next to in_axes in data pipelines","Remember nested vmap axes refer to the original rank"],"tags":["jax","vmap","shape","rank","in-axes"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}