jax-ml/jax · error · ValueError

{name} was requested to map a value of non-array type {core.

Error message

{name} was requested to map a value of non-array type {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.

What it means

Raised when vmap/pmap/linearize tries to map a non-array JAX type (e.g. a string, dict, or custom high-level type) along an integer in_axes entry. Non-array pytree nodes cannot have an integer axis mapped; JAX requires a MappingSpec and an explicit axis_size instead.

Source

Thrown at jax/_src/api.py:1316

    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 "
            "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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set in_axes=None for the non-array argument
2. Pass a MappingSpec instance for that argument's in_axes entry and axis_size explicitly
3. Convert the metadata to static/closure state or a JAX array

Example fix

// before
jax.vmap(f, in_axes=(0, 0))(x, meta_dict)
// after
jax.vmap(f, in_axes=(0, None))(x, meta_dict)

for leaf, ax in zip(tree_leaves(args), tree_leaves(in_axes)):
    if isinstance(ax, int) and not (hasattr(leaf, 'ndim') or core.valid_jaxtype(leaf) and not core.typeof(leaf).is_high):
        raise TypeError('non-array leaf needs in_axes=None or MappingSpec')

def is_mappable_array(x): return hasattr(x, 'shape') and hasattr(x, 'dtype')

When it happens

Trigger: jax.vmap(f)(x, {'a': 1}) with in_axes=0 where the dict argument is a non-array leaf; mapping a custom type with a plain integer in_axes.

Common situations: Passing dicts of metadata or custom vmappable objects through vmap with default in_axes=0; new JAX versions that enforce MappingSpec for high-level types.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• {name} was requested to map its argument along axis {axis},
• {name} must have at least one non-None value in in_axes or a
• Currently only support batch_dim in [0, None], but got {dim=
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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