jax-ml/jax · error · ValueError

{name} wrapped function must be passed at least one argument

Error message

{name} wrapped function must be passed at least one argument containing an array or axis_size must be specified, got empty *args={args} and **kwargs={kwargs}

What it means

Raised by vmap/pmap-style transforms when the wrapped function receives no arguments containing arrays (or any mapped leaves) and no axis_size was given, so the mapped axis size cannot be inferred.

Source

Thrown at jax/_src/api.py:1301

def _check_ema_unmapped_args(ema, args_flat, in_axes_flat):
  if ema is None:
    return
  for a, i in zip(args_flat, in_axes_flat):
    if i is None:
      aval = core.typeof(a)
      spec = set(sharding_impls.flatten_spec(aval.sharding.spec))
      if any(e in spec for e in ema):
        raise ValueError(
            "Unmapped values passed to vmap cannot be sharded along the mesh"
            f" axis you are vmapping over. Got type: {aval.str_short(True)},"
            f" in_axes: {i} and vmapped mesh axis: {ema}")

def _mapped_axis_size(fn, tree, vals, dims, name, axis_size=None):
  if not vals:
    if axis_size is not None:
      return axis_size
    args, kwargs = tree_unflatten(tree, vals)
    raise ValueError(
        f"{name} wrapped function must be passed at least one argument "
        "containing an array or axis_size must be specified, got empty "
        f"*args={args} and **kwargs={kwargs}"
    )

  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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass axis_size=N to vmap so the mapped size is explicit
Add at least one array argument to the function and map it via in_axes
Move constants into the function's arguments as arrays

Example fix

// before
jax.vmap(lambda: jnp.arange(3) + 1)()
// after
jax.vmap(lambda n: jnp.arange(3) + n, axis_size=8)()
# or
jax.vmap(lambda n: jnp.arange(3) + n)(jnp.ones(8))

Defensive patterns

Strategy: validation

Validate before calling

if not any(hasattr(l, 'ndim') for l in tree_leaves((args, kwargs))):
    assert axis_size is not None, 'pass axis_size when no array args are mapped'

Prevention

Always pass axis_size for vmap over zero-array functions
Prefer array arguments over closures for mapped values
Wrap risky calls: jax.vmap(f, axis_size=N) by default in such utilities

When it happens

Trigger: jax.vjp-like call such as jax.vmap(lambda: ... )() or jax.vmap(f)() where all arguments are empty/pytrees with no leaves, or jax.linearize/jax.vjp on a zero-argument function, without passing axis_size.

Common situations: Refactoring a function to take configuration via closure instead of arguments; calling vmap over a function of only Python scalars/static values; testing with dummy empty inputs.

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

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