{"record":{"id":"957c3eb4e4cfb6f0","repo":"jax-ml/jax","slug":"linearized-function-called-on-tangent-values-incon","errorCode":null,"errorMessage":"linearized function called on tangent values inconsistent with the original primal values:\nGot tangent aval {tangent_aval} for primal aval {primal_aval} but expected {expected_tangent_aval}.{extra_msg}","messagePattern":"linearized function called on tangent values inconsistent with the original primal values:\nGot tangent aval (.+?) for primal aval (.+?) but expected (.+?)\\.(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1602,"sourceCode":"    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\"\n              \" the tangent value passed to the callable `f_jvp` returned by\"\n              \" `jax.linearize`\")\n        extra_msg = \" \\nThis might be fixed by:\\n\" + \"\\n\".join(\n            f\"  * {d};\" for d in pvary_applications)\n      raise ValueError(\n          \"linearized function called on tangent values inconsistent with \"\n          \"the original primal values:\\n\"\n          f\"Got tangent aval {tangent_aval} for primal aval {primal_aval} \"\n          f\"but expected {expected_tangent_aval}.{extra_msg}\")\n  sres_flat = tree_leaves(structured_residuals)\n  tangents_out = eval_jaxpr(jaxpr, consts, *tangents_ft, *sres_flat)\n  tangents_out_ = iter(tangents_out)\n  full_out = [a2tz(aval).instantiate() if known else next(tangents_out_)\n              for aval, known in zip(out_avals, out_zeros)]\n  assert next(tangents_out_, None) is None\n  return out_avals.update(full_out).unflatten()\n\n# TODO(mattjj): see similar function in custom_derivatives.py\ndef _temporary_dtype_exception(a, a_) -> bool:\n  if isinstance(a, core.ShapedArray) and isinstance(a_, core.ShapedArray):\n    return a.shape == a_.shape and a_.dtype == float0\n  return False\n","sourceCodeStart":1584,"sourceCodeEnd":1620,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1584-L1620","documentation":"The linearized callable checks each tangent leaf's abstract value (shape/dtype, including varying/precision qualifiers) against the tangent aval expected from the original primal. Mismatches raise this, with suggestions (e.g. applying jax.lax.pcast with to='varying') when relevant.","triggerScenarios":"Calling jax.linearize(f)(...) result with tangents of wrong dtype/shape, or with non-'varying' values where the primal was varying-precision and pcast is required on tangents.","commonSituations":"Mixed-precision (varying) pipelines where tangents must be marked varying via jax.lax.pcast; changing precision context between linearize and the later call; reusing tangents from a differently-shaped input.","solutions":["Follow the extra_msg hints: apply jax.lax.pcast(t, to='varying') to tangent values when primals are varying","Rebuild tangents from the exact primals used at linearize time (zeros_like/instantiate)","Ensure shape/dtype/precision context matches between linearize and the call"],"exampleFix":"# before\nf_lin = jax.linearize(f)(x)\ny = f_lin(jnp.ones_like(x))\n# after (varying-precision primals)\ny = f_lin(jax.lax.pcast(jnp.ones_like(x), (), to='varying'))","handlingStrategy":"validation","validationCode":"for pa, t in zip(in_avals, tree_leaves(tangents)):\n    exp = pa.to_tangent_aval()\n    assert core.typecompat(exp, jax.core.typeof(t)), f'{t} incompatible with {exp}'","typeGuard":"def tangent_aval_ok(primal_aval, t): return core.typecompat(primal_aval.to_tangent_aval(), jax.core.typeof(t))","tryCatchPattern":null,"preventionTips":["Apply jax.lax.pcast(..., to='varying') to tangents of varying primals","Keep precision context identical at linearize and call time","Validate tangents against primal avals in wrappers"],"tags":["jax","linearize","tangent","aval","mixed-precision"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}