{"record":{"id":"d07b5bb92cbc7de5","repo":"jax-ml/jax","slug":"expected-in-avals-tree-got-tangents-ft-tree","errorCode":null,"errorMessage":"expected {in_avals.tree}, got {tangents_ft.tree}","messagePattern":"expected (.+?), got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1579,"sourceCode":"  primals_ft = ft.flatten(primals)\n  in_nzs_flat = None if in_nzs is None else tuptree_flags(\n      in_nzs, primals_ft.tree, 'in_nzs', 'the in_nzs argument to jax.linearize')\n  out_primals_ft, out_zeros, jaxpr, consts, structured_residuals, *maybe_aux = \\\n      ad.linearize(fun, primals_ft, has_aux=has_aux, in_nzs=in_nzs_flat)\n  in_avals = primals_ft.map(core.typeof)\n  out_avals = out_primals_ft.map(core.typeof)\n  lifted_jvp = Partial(\n      partial(_lift_linearized, jaxpr, in_avals, out_avals, out_zeros),\n      consts, structured_residuals)\n  lifted_jvp.out_nzs = tuple(not z for z in out_zeros)  # pyrefly: ignore[missing-attribute]\n  return out_primals_ft.unflatten(), lifted_jvp, *maybe_aux\n\n\ndef _lift_linearized(jaxpr, in_avals, out_avals, out_zeros, consts,\n                     structured_residuals, *tangents):\n  tangents_ft = ft.flatten(tangents)\n  if tangents_ft.tree != in_avals.tree:\n    raise TypeError(f\"expected {in_avals.tree}, got {tangents_ft.tree}\")\n\n  tangent_avals = tangents_ft.map(core.typeof)\n  for primal_aval, tangent_aval in zip(in_avals, tangent_avals):\n    expected_tangent_aval  = primal_aval.to_tangent_aval()\n    if not core.typecompat(expected_tangent_aval, tangent_aval):\n      extra_msg = ''\n      if (isinstance(primal_aval, core.ShapedArray) and\n          isinstance(tangent_aval, core.ShapedArray) and\n          primal_aval.mat != tangent_aval.mat):\n        # TODO(yashkatariya): Tweak error.\n        pvary_applications = []\n        if left := tangent_aval.mat.varying - primal_aval.mat.varying:\n          pvary_applications.append(\n              f\"applying `jax.lax.pcast(..., {tuple(left)}, to='varying')` to\"\n              \" the primal value passed to `jax.linearize`\")\n        if left := primal_aval.mat.varying - tangent_aval.mat.varying:\n          pvary_applications.append(\n              f\"applying `jax.lax.pcast(..., {tuple(left)}, to='varying')` to\"","sourceCodeStart":1561,"sourceCodeEnd":1597,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1561-L1597","documentation":"When calling the function returned by jax.linearize, the tangents passed in must match the pytree structure of the original primal arguments. This TypeError reports the expected vs got tree structures.","triggerScenarios":"linearize(f)(x)(t) where t has different structure than x (e.g. x was a tuple (a, b) but a flat tangent or differently nested structure is passed).","commonSituations":"Saving/reloading a linearized function and calling it with differently-packed inputs; partial-application refactor changing arity; passing cotangents from a vjp (output structure) into a linearized call (input structure).","solutions":["Pass tangents structured exactly like the original primals (reuse the primal pytree via tree.map)","Capture the original tree and rebuild tangents with jax.tree.unflatten before calling","Write a small wrapper that flattens inputs the same way as at linearize time"],"exampleFix":"# before\nf_lin = jax.linearize(f)(x, y)\nout = f_lin(jnp.ones_like(x))  # missing y\n# after\nf_lin = jax.linearize(f)(x, y)\nout = f_lin(jnp.ones_like(x), jnp.zeros_like(y))","handlingStrategy":"validation","validationCode":"assert jax.tree.structure(tangents) == jax.tree.structure(original_primals), 'linearize tangents must mirror primal tree'","typeGuard":"def tangent_tree_ok(primals, t): return jax.tree.structure(t) == jax.tree.structure(primals)","tryCatchPattern":null,"preventionTips":["Store the primal tree with the linearized function and rebuild inputs via unflatten","Call linearized callables with tree.map over the original primals","Cover linearize call signatures in unit tests"],"tags":["jax","linearize","jvp","pytree"],"backgroundTag":"structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}