{"record":{"id":"0fbcda01de3a8d02","repo":"jax-ml/jax","slug":"unexpected-jax-type-e-g-shape-dtype-for-argumen","errorCode":null,"errorMessage":"unexpected JAX type (e.g. shape/dtype) for argument to VJP function: got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} because the corresponding output of the differentiated function had JAX type {aval.str_short()}","messagePattern":"unexpected JAX type \\(e\\.g\\. shape/dtype\\) for argument to VJP function: got (.+?), but expected (.+?) because the corresponding output of the differentiated function had JAX type (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/api.py","lineNumber":1977,"sourceCode":"\nBut the tree structures differ:\n\"\"\"\n  msg += '\\n'.join(f\"  * out{keystr(path)} was a {thing1} in the original \"\n                   f\"output, but a {thing2} here, so {explanation}.\"\n                   for path, thing1, thing2, explanation\n                   in equality_errors_pytreedef(out_tree, ct_tree))\n  raise ValueError(msg)\n\n\ndef _vjp_check_ct_avals(cts, primal_avals):\n  # TODO(mattjj): improve this error  by flattening with keys in the first place\n  for ct, aval in zip(cts, primal_avals):\n    if isinstance(ct, ad.Zero): continue\n    ct_aval = typeof(ct)\n    ct_aval_expected = aval.to_ct_aval()\n    if (not core.typecompat(ct_aval, ct_aval_expected) and\n        not _temporary_dtype_exception(ct_aval, ct_aval_expected)):\n      raise ValueError(\n          \"unexpected JAX type (e.g. shape/dtype) for argument to VJP function: \"\n          f\"got {ct_aval.str_short()}, but expected {ct_aval_expected.str_short()} \"\n          \"because the corresponding output of the differentiated function had JAX type \"\n          f\"{aval.str_short()}\")\n\n\n@register_dataclass\n@dataclasses.dataclass(frozen=True, slots=True)\nclass NotNeeded:\n  pass\n\n@register_dataclass\n@dataclasses.dataclass(frozen=True, slots=True)\nclass NotSaveable:\n  pass\n\n@dataclasses.dataclass(frozen=True, slots=True)\nclass GradValue:","sourceCodeStart":1959,"sourceCodeEnd":1995,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/api.py#L1959-L1995","documentation":"In jax.vjp, each cotangent passed to the backward function must have the same shape and dtype (JAX aval) as the corresponding primal output's cotangent aval. JAX raises this when e.g. the cotangent is float64 while the output was float32, or has a different shape.","triggerScenarios":"Passing a numpy float64 array as cotangent for a float32 output; passing a cotangent with wrong shape (e.g. scalar for a (3,) output); mixing jnp and np scalars with different default dtypes.","commonSituations":"x64 mode enabled (jax.config.update('jax_enable_x64', True)) so np.ones gives float64 while the traced output is float32; constructing cotangents from numpy instead of jnp.zeros_like(out); shape errors from broadcasting assumptions.","solutions":["Build cotangents with jax.tree_util.tree_map(jnp.zeros_like, out) or jnp.ones_like","Match dtypes explicitly: jnp.ones(shape, dtype=out.dtype)","Check jax_enable_x64 consistency between how the function was traced and how cotangents are constructed"],"exampleFix":"// before\ncts = (np.ones(3.0),)   # float64 vs float32 output\n// after\ncts = (jnp.ones_like(out_array),)","handlingStrategy":"validation","validationCode":"import jax.tree_util as jtu\ncts = jtu.tree_map(lambda o: jnp.zeros(o.shape, o.dtype), out)","typeGuard":null,"tryCatchPattern":"try:\n    in_cts = vjp_fn(cts)\nexcept ValueError as e:\n    if 'unexpected JAX type' in str(e):\n        cts = jtu.tree_map(lambda o: jnp.zeros_like(o), out)\n        in_cts = vjp_fn(cts)\n    else:\n        raise","preventionTips":["Never build cotangents with raw numpy when outputs are traced jax arrays","Pin dtypes with out.dtype rather than relying on defaults"],"tags":["jax","vjp","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"}