jax-ml/jax · error · ValueError

{name} was requested to map its argument along axis {axis},

Error message

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

What it means

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.

Source

Thrown at jax/_src/api.py:1324

  def _get_axis_size(name: str, x, axis: int) -> core.AxisSize | None:
    shape: tuple[core.AxisSize, ...] = ()
    try:
      shape = np.shape(x)
      return shape[axis]
    except (IndexError, TypeError) as e:
      if not core.valid_jaxtype(x) or not isinstance(axis, int):
        return None  # Suppress the check for custom vmappable types.
      if core.typeof(x).is_high:
        raise ValueError(
            f"{name} was requested to map a value of non-array type "
            f"{core.typeof(x)} along axis {axis}, but non-array types can't "
            "be mapped along an integer axis. Instead pass a mapping spec (a "
            "MappingSpec instance) as this argument's in_axes entry, and "
            "pass axis_size explicitly.") from None
      min_rank = axis + 1 if axis >= 0 else -axis
      # TODO(mattjj): better error message here
      raise ValueError(
          f"{name} was requested to map its argument along axis {axis}, "
          f"which implies that its rank should be at least {min_rank}, "
          f"but is only {len(shape)} (its shape is {shape})") from e

  all_mapped_sizes = [
    None if d is None else _get_axis_size(name, x, d)
    for x, d in zip(vals, dims)
  ]
  all_sizes = [s for s in all_mapped_sizes if s is not None]
  if axis_size is not None:
    all_sizes.append(axis_size)
  sizes = core.dedup_referents(all_sizes)
  if len(sizes) == 1:
    sz, = sizes
    return sz
  if not sizes:
    raise ValueError(f"{name} must have at least one non-None value in in_axes "
                     "or axis_size must be specified")

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
jax.vmap(f, in_axes=1)(jnp.zeros(4))
// after
jax.vmap(f, in_axes=0)(jnp.zeros((4, 4)))
# or
jax.vmap(f, in_axes=1)(jnp.zeros(4)[:, None])

Defensive patterns

Strategy: validation

Validate before calling

for leaf, ax in zip(tree_leaves(args), tree_leaves(in_axes)):
    if isinstance(ax, int) and hasattr(leaf, 'ndim'):
        assert -leaf.ndim <= ax < leaf.ndim, f'in_axes {ax} invalid for shape {leaf.shape}'

Type guard

def axis_valid(x, ax): return hasattr(x, 'ndim') and -x.ndim <= ax < x.ndim

Prevention

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

When it happens

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

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

Related errors

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