jax-ml/jax · error · ValueError

{name} must have at least one non-None value in in_axes or a

Error message

{name} must have at least one non-None value in in_axes or axis_size must be specified

What it means

Raised when all in_axes entries are None (nothing is actually mapped) and no axis_size was provided, so vmap cannot determine the size of the mapped axis. vmap requires at least one mapped value or an explicit axis_size.

Source

Thrown at jax/_src/api.py:1341

      # 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")

  def _get_argument_type(x):
    try:
      return shaped_abstractify(x).str_short()
    except TypeError: # Catch all for user specified objects that can't be interpreted as a data type
      return "unknown"
  msg = [f"{name} got inconsistent sizes for array axes to be mapped:\n"]
  args, kwargs = tree_unflatten(tree, vals)
  try:
    ba = inspect.signature(fn).bind(*args, **kwargs)
    signature_parameters: list[str] | None = list(ba.signature.parameters.keys())
  except (TypeError, ValueError):
    signature_parameters = None

  def arg_name(key_path):
    if signature_parameters is None:
      return f"args{keystr(key_path)}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass axis_size=N explicitly
2. Set at least one in_axes entry to an integer axis of a corresponding array argument
3. Double-check that you actually want vmap if nothing is mapped

Example fix

// before
jax.vmap(f, in_axes=None)(x)
// after
jax.vmap(f, in_axes=None, axis_size=16)(x)
# or
jax.vmap(f, in_axes=0)(x)

if all(a is None for a in tree_leaves(in_axes)):
    assert axis_size is not None, 'vmap with all-None in_axes needs axis_size'

When it happens

Trigger: jax.vmap(f, in_axes=None)(x); jax.vmap(f, in_axes=(None, None))(x, y) without axis_size.

Common situations: Dynamically computed in_axes that end up all None; migrating code where the batch axis was removed; wrapping functions that only return constants.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

• {name} wrapped function must be passed at least one argument
• {name} was requested to map a value of non-array type {core.
• {name} was requested to map its argument along axis {axis},
• Currently only support batch_dim in [0, None], but got {dim=
• pmap requires at least one argument with a mapped axis.

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