{"record":{"id":"0c59b8e82d27ffbc","repo":"jax-ml/jax","slug":"cotangent-type-does-not-match-function-output-exp","errorCode":null,"errorMessage":"cotangent type does not match function output, expected {out_avals} but got {out_cts}","messagePattern":"cotangent type does not match function output, expected (.+?) but got (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":2124,"sourceCode":"                                                      instantiate=True)\n  jaxpr, _ = pe.dce_jaxpr(jaxpr, [True] * len(jaxpr.outvars), True)\n  out_avals, _ = unzip2(out_pvals)\n  out_dtypes = [a.dtype for a in out_avals if not a.is_high]\n  if not (all(dtypes.issubdtype(d, np.inexact) for d in in_dtypes + out_dtypes)\n          or all(dtypes.issubdtype(d, np.integer)\n                 for d in in_dtypes + out_dtypes)):\n    raise TypeError(\"linear_transpose only supports [float or complex] -> \"\n                    \"[float or complex], and integer -> integer functions, \"\n                    f\"but got {in_dtypes} -> {out_dtypes}.\")\n\n  @api_boundary\n  def transposed_fun(const, out_cotangent):\n    out_cts, out_tree2 = tree_flatten(out_cotangent)\n    if out_tree() != out_tree2:\n      raise TypeError(\"cotangent tree does not match function output, \"\n                      f\"expected {out_tree()} but got {out_tree2}\")\n    if not all(map(core.typecheck, out_avals, out_cts)):\n      raise TypeError(\"cotangent type does not match function output, \"\n                      f\"expected {out_avals} but got {out_cts}\")\n    dummies = [ad.UndefinedPrimal(a.to_ct_aval()) for a in in_avals]\n    in_cts = ad.backward_pass(jaxpr, True, const, dummies, out_cts)\n    in_cts = map(ad.instantiate_zeros, in_cts)\n    return tree_unflatten(in_tree, in_cts)\n\n  # Ensure that transposed_fun is a PyTree\n  return Partial(transposed_fun, const)\n\n\n@overload\ndef make_jaxpr(\n    fun: Callable,\n    static_argnums: int | Sequence[int] = (),\n    axis_env: Sequence[tuple[AxisName, int]] | None = None,\n    return_shape: Literal[False] = ...,\n) -> Callable[..., core.Jaxpr]:\n  ...","sourceCodeStart":2106,"sourceCodeEnd":2142,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L2106-L2142","documentation":"Even when the pytree structure matches, each cotangent leaf passed to a jax.linear_transpose pullback must typecheck against the corresponding output aval (shape and dtype). This fires on shape or dtype mismatches per leaf.","triggerScenarios":"Passing a float64 cotangent for a float32 output; passing shape (2,3) where output was (3,2); passing a Python scalar where a tracer-compatible array is required.","commonSituations":"x64 enabled causing numpy-built cotangents to be float64; reshaping/broadcast mistakes; passing zeros with an inferred dtype that differs from the output dtype.","solutions":["Construct cotangents with jnp.ones_like / jnp.zeros_like on the actual outputs","Specify dtype explicitly: jnp.zeros(shape, dtype=out.dtype)","Ensure jax_enable_x64 setting matches between tracing and cotangent construction"],"exampleFix":"# before\nct = np.zeros(x.shape)          # possibly float64\n# after\nct = jnp.zeros_like(f(x))","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\nout = f(*primals)\nct = jtu.tree_map(lambda o: jnp.zeros(o.shape, o.dtype), out)  # dtype/shape-exact","typeGuard":null,"tryCatchPattern":"try:\n    in_ct = f_t(ct)\nexcept TypeError as e:\n    if 'cotangent type' in str(e):\n        ct = jtu.tree_map(jnp.zeros_like, f(*primals))\n        in_ct = f_t(ct)\n    else:\n        raise","preventionTips":["Build cotangents with jnp.zeros_like/ones_like on real outputs","Keep jax_enable_x64 consistent across your program"],"tags":["jax","linear-transpose","dtype","shape-mismatch"],"backgroundTag":"cotangent-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}