{"record":{"id":"97467fa1dc2a4ba7","repo":"jax-ml/jax","slug":"join-msg-2","errorCode":null,"errorMessage":"{''.join(msg)[:-2]}","messagePattern":"\\{''\\.join\\(msg\\)\\[:-2\\]\\}","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1401,"sourceCode":"    for i, isz in enumerate(all_mapped_sizes):\n      if core.definitely_equal(isz, sz): return i\n    assert False, (sz, all_mapped_sizes)\n\n  ex, *examples = (key_paths[_all_sizes_index(sz)] for sz, _ in counts)\n  ax, *axs = (dims[_all_sizes_index(sz)] for sz, _ in counts)\n\n  if axis_size is not None:\n    msg.append(f\"  * the `axis_size` argument was {axis_size};\\n\")\n  if ct == 1:\n    msg.append(f\"  * one axis had size {sz}: axis {ax} of {ex};\\n\")\n  else:\n    msg.append(f\"  * most axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\\n\")\n  for ex, ax, (sz, ct) in zip(examples, axs, other_counts):\n    if ct == 1:\n      msg.append(f\"  * one axis had size {sz}: axis {ax} of {ex};\\n\")\n    else:\n      msg.append(f\"  * some axes ({ct} of them) had size {sz}, e.g. axis {ax} of {ex};\\n\")\n  raise ValueError(''.join(msg)[:-2])  # remove last semicolon and newline\n\n\n@partial(api_boundary, repro_api_name=\"jax.jvp\")\ndef jvp(\n    fun: Callable, primals, tangents, has_aux: bool = False\n  ) -> tuple[Any, ...]:\n  \"\"\"Computes a (forward-mode) Jacobian-vector product of ``fun``.\n\n  Args:\n    fun: Function to be differentiated. Its arguments should be arrays, scalars,\n      or standard Python containers of arrays or scalars. It should return an\n      array, scalar, or standard Python container of arrays or scalars.\n    primals: The primal values at which the Jacobian of ``fun`` should be\n      evaluated. Should be either a tuple or a list of arguments,\n      and its length should be equal to the number of positional parameters of\n      ``fun``.\n    tangents: The tangent vector for which the Jacobian-vector product should be\n      evaluated. Should be either a tuple or a list of tangents, with the same","sourceCodeStart":1383,"sourceCodeEnd":1419,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1383-L1419","documentation":"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.","triggerScenarios":"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].","commonSituations":"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.","solutions":["Make the mapped dimensions of all arguments equal (fix data shapes or batching logic)","Correct the in_axes so each argument's mapped axis refers to the shared batch dimension","Print shapes of all vmap arguments before the call to find the odd one out"],"exampleFix":"// before\njax.vmap(lambda a, b: a + b)(jnp.zeros(32), jnp.zeros(64))\n// after\nb = jnp.zeros((32, 64))\njax.vmap(lambda a, b: a + b)(jnp.zeros(32), b)  # map axis 0 of both","handlingStrategy":"validation","validationCode":"sizes = {np.shape(l)[d] for l, d in zip(tree_leaves(args), tree_leaves(in_axes)) if d is not None}\nassert len(sizes) <= 1, f'mapped size mismatch: {sizes}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check set of mapped-axis sizes before vmap","Keep one authoritative batch_size variable in the data pipeline","Print all arg shapes when debugging vmap size errors"],"tags":["jax","vmap","shape","batch-size"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}