jax-ml/jax · error · ValueError

{''.join(msg)[:-2]}

Error message

{''.join(msg)[:-2]}

What it means

Raised by vmap's axis-size deduction when the mapped arguments have inconsistent sizes along their mapped axes (e.g. one has batch size 32 and another 64). The message enumerates how many axes had each size with examples so the mismatch is easy to locate.

Source

Thrown at jax/_src/api.py:1401

    for i, isz in enumerate(all_mapped_sizes):
      if core.definitely_equal(isz, sz): return i
    assert False, (sz, all_mapped_sizes)

  ex, *examples = (key_paths[_all_sizes_index(sz)] for sz, _ in counts)
  ax, *axs = (dims[_all_sizes_index(sz)] for sz, _ in counts)

  if axis_size is not None:
    msg.append(f"  * the `axis_size` argument was {axis_size};\n")
  if ct == 1:
    msg.append(f"  * one axis had size {sz}: axis {ax} of {ex};\n")
  else:
    msg.append(f"  * most axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n")
  for ex, ax, (sz, ct) in zip(examples, axs, other_counts):
    if ct == 1:
      msg.append(f"  * one axis had size {sz}: axis {ax} of {ex};\n")
    else:
      msg.append(f"  * some axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\n")
  raise ValueError(''.join(msg)[:-2])  # remove last semicolon and newline


@partial(api_boundary, repro_api_name="jax.jvp")
def jvp(
    fun: Callable, primals, tangents, has_aux: bool = False
  ) -> tuple[Any, ...]:
  """Computes a (forward-mode) Jacobian-vector product of ``fun``.

  Args:
    fun: Function to be differentiated. Its arguments should be arrays, scalars,
      or standard Python containers of arrays or scalars. It should return an
      array, scalar, or standard Python container of arrays or scalars.
    primals: The primal values at which the Jacobian of ``fun`` should be
      evaluated. Should be either a tuple or a list of arguments,
      and its length should be equal to the number of positional parameters of
      ``fun``.
    tangents: The tangent vector for which the Jacobian-vector product should be
      evaluated. Should be either a tuple or a list of tangents, with the same

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the mapped dimensions of all arguments equal (fix data shapes or batching logic)
2. Correct the in_axes so each argument's mapped axis refers to the shared batch dimension
3. Print shapes of all vmap arguments before the call to find the odd one out

Example fix

// before
jax.vmap(lambda a, b: a + b)(jnp.zeros(32), jnp.zeros(64))
// after
b = jnp.zeros((32, 64))
jax.vmap(lambda a, b: a + b)(jnp.zeros(32), b)  # map axis 0 of both

sizes = {np.shape(l)[d] for l, d in zip(tree_leaves(args), tree_leaves(in_axes)) if d is not None}
assert len(sizes) <= 1, f'mapped size mismatch: {sizes}'

When it happens

Trigger: jax.vmap(f)(jnp.zeros(32), jnp.zeros(64)); jax.vmap(f, in_axes=(0, 1))(a, b) where a.shape[0] != b.shape[1].

Common situations: Off-by-one data loading, misaligned batch dimension (using axis 0 of one array and axis 1 of another with different lengths), padding/trimming bugs in a preprocessing step.

Related errors

• {name} was requested to map its argument along axis {axis},
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• {name} wrapped function must be passed at least one argument
• {name} was requested to map a value of non-array type {core.

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