{"record":{"id":"858013c191e35efa","repo":"jax-ml/jax","slug":"jvp-called-with-different-primal-and-tangent-shape","errorCode":null,"errorMessage":"jvp called with different primal and tangent shapes;Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}","messagePattern":"jvp called with different primal and tangent shapes;Got primal shape (.+?) and tangent shape as (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1469,"sourceCode":"\ndef _jvp(fun: Callable, primals, tangents, has_aux=False):\n  ps_ft = ft.flatten(primals)\n  ts_ft = ft.flatten(tangents)\n  if ps_ft.tree != ts_ft.tree:\n    raise TypeError(\"primal and tangent arguments to jax.jvp must have the same tree \"\n                    f\"structure; primals have tree structure {ps_ft.tree} whereas tangents have \"\n                    f\"tree structure {ts_ft.tree}.\")\n  for p, t in zip(ps_ft, ts_ft):\n    if not isinstance(core.typeof(p), ShapedArray): continue\n    if core.primal_dtype_to_tangent_dtype(_dtype(p)) != _dtype(t):\n      raise TypeError(\"primal and tangent arguments to jax.jvp do not match; \"\n                      \"dtypes must be equal, or in case of int/bool primal dtype \"\n                      \"the tangent dtype must be float0.\"\n                      f\"Got primal dtype {_dtype(p)} and so expected tangent dtype \"\n                      f\"{core.primal_dtype_to_tangent_dtype(_dtype(p))}, but got \"\n                      f\"tangent dtype {_dtype(t)} instead.\")\n    if np.shape(p) != np.shape(t):\n      raise ValueError(\"jvp called with different primal and tangent shapes;\"\n                       f\"Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}\")\n\n  out_primals, out_tangents, *aux = ad.jvp(fun, ps_ft, ts_ft, has_aux=has_aux)\n  return out_primals.unflatten(), out_tangents.unflatten(), *aux\n\n@overload\ndef linearize(fun: Callable, *primals, has_aux: Literal[False] = False,\n              in_nzs: Any = None) -> tuple[Any, Callable]:\n  ...\n\n@overload\ndef linearize(fun: Callable, *primals, has_aux: Literal[True],\n              in_nzs: Any = None) -> tuple[Any, Callable, Any]:\n  ...\n\n@partial(api_boundary, repro_api_name=\"jax.linearize\")\ndef linearize(fun: Callable, *primals, has_aux: bool = False,\n              in_nzs: Any = None","sourceCodeStart":1451,"sourceCodeEnd":1487,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1451-L1487","documentation":"Within jax.jvp, each primal leaf and its corresponding tangent leaf must have identical shapes. This error reports both shapes when they differ (dtype/tree checks already passed).","triggerScenarios":"jax.jvp(f, (jnp.zeros((2,3)),), (jnp.zeros((3,2)),)); tangents built from a differently-shaped array (e.g. transposed or reshaped) than primals.","commonSituations":"Tangent constructed from a different variable than its primal after refactors; broadcasting assumptions where shapes happen to differ in one dimension; transposed weights in custom gradients.","solutions":["Generate tangents via jax.tree.map(jnp.zeros_like, primals) so shapes match by construction","Fix the tangent construction to use the corresponding primal's shape","Add a shape equality assert before the jvp call"],"exampleFix":"# before\njax.jvp(f, (w,), (jnp.zeros_like(w.T),))\n# after\njax.jvp(f, (w,), (jnp.zeros_like(w),))","handlingStrategy":"validation","validationCode":"for p, t in zip(tree_leaves(primals), tree_leaves(tangents)):\n    if hasattr(p, 'shape'):\n        assert np.shape(p) == np.shape(t), f'shape mismatch {np.shape(p)} vs {np.shape(t)}'","typeGuard":"def shapes_match(p, t): return not hasattr(p, 'shape') or np.shape(p) == np.shape(t)","tryCatchPattern":null,"preventionTips":["Generate tangents with zeros_like of the exact primal","Avoid transposed/reshaped buffers as tangents","Assert shape equality in tests covering jvp paths"],"tags":["jax","jvp","shape","autodiff"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}