{"record":{"id":"b53217cb5b93d4c5","repo":"jax-ml/jax","slug":"at-keystr-path-got-fwd-output-type-ty-str-sho","errorCode":null,"errorMessage":"at {keystr(path)}, got fwd output type {ty.str_short()} which doesn't match primal output type {primal_aval.str_short()}","messagePattern":"at (.+?), got fwd output type (.+?) which doesn't match primal output type (.+?)","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":890,"sourceCode":"\ndef _vjp_primal_fwd_tree_mismatch_err(self, tree):\n  return (f\"Custom VJP fwd rule {self.fwd.__name__} for function {self.traced.fun_name} \"\n          \"must produce a pair (list or tuple of length two) where the first \"\n          \"element represents the primal output \"\n          \"(equal to the output of the custom_vjp-decorated function \"\n          f\"{self.traced.fun_name}) and the \"\n          \"second element represents residuals (i.e. values stored from the \"\n          \"forward pass for use on the backward pass), but \"\n          f\"instead the fwd rule output's first element had container/pytree \"\n          \"structure:\\n\"\n          f\"\"\"    {str(tree ).replace(\"'\", \"\")}\\n\"\"\"\n          f\"while the custom_vjp-decorated function {self.traced.fun_name} had output \"\n          \"container/pytree structure:\\n\"\n          f\"\"\"    {str(self.out_tree).replace(\"'\", \"\")}.\"\"\")\n\ndef _vjp_fwd_aval_mismatch_err(path, primal_aval, fwd_val):\n  if not core.typematch(ty := typeof(fwd_val), primal_aval):\n    raise TypeError(f\"at {keystr(path)}, got fwd output type {ty.str_short()} \"\n                    f\"which doesn't match primal output type {primal_aval.str_short()}\")\n\ndef _vjp_bwd_aval_mismatch_err(primal_sourceinfo, path, primal_aval, ct):\n  if config.disable_bwd_checks.value:\n    return\n  if isinstance(ct, ad_util.Zero):\n    return\n  if isinstance(primal_aval, AbstractRef):\n    primal_aval = primal_aval.inner_aval\n  expected = primal_aval.to_ct_aval()\n  ct_aval = ct.aval if isinstance(ct, ad_util.SymbolicZero) else typeof(ct)\n  if (not core.typematch(expected, ct_aval) and\n      not _temporary_dtype_exception(expected, ct_aval) and\n      getattr(expected, 'dtype', None) is not dtypes.float0):\n    result = f\"at output{keystr(path)} \" if path else \"\"\n    raise ValueError(\n        f\"{result}the bwd rule attached to {primal_sourceinfo} produced an\"\n        f\" output of type {ct_aval.str_short()} which doesn't match expected\"","sourceCodeStart":872,"sourceCodeEnd":908,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L872-L908","documentation":"During custom_vjp evaluation, JAX checks that each output of the user-supplied forward function has the same type/aval as the corresponding primal output of the @jax.custom_vjp-decorated function. If the fwd rule returns a value whose dtype/shape/sharding differs from the primal's, this TypeError is raised at the mismatching pytree path.","triggerScenarios":"Passing a fwd function to custom_vjp.defvjp whose return values (e.g. residuals or primals) differ in shape or dtype from what the decorated function returns, e.g. returning float32 residuals when the primal outputs float64, or returning a differently nested pytree.","commonSituations":"Writing a custom_vjp fwd rule that casts or reshapes outputs (e.g. l.to(torch.float32) style casts, or int truncation) inconsistently with the primal function; changing the primal signature without updating the fwd rule; enabling x64 mode after writing rules with hardcoded dtypes.","solutions":["Make the fwd rule return the primal outputs with exactly the same avals (shape+dtype) as the decorated function","Check for implicit dtype casts in the fwd function (e.g. Python floats mixing with float32, x64 mismatches) and cast explicitly with jax.lax.convert_element_type to the primal dtype","Ensure the pytree structure returned by fwd matches the decorated function's output structure exactly"],"exampleFix":"// before\n@jax.custom_vjp\ndef f(x):\n  return x.astype(jnp.float32) * 2\ndef f_fwd(x):\n  return x * 2, None          # returns float64 if x is float64\n// after\ndef f_fwd(x):\n  primal = x.astype(jnp.float32) * 2   # match primal aval exactly\n  return primal, None","handlingStrategy":"validation","validationCode":"# verify fwd outputs match primal avals before use\nimport jax\nflat_primal = jax.eval_shape(f, x)\nflat_fwd, _ = jax.eval_shape(lambda x: (f_fwd(x)), x)  # first of pair\nassert jax.tree_util.tree_structure(flat_fwd) == jax.tree_util.tree_structure(flat_primal)","typeGuard":null,"tryCatchPattern":"try:\n    y = f(x)\nexcept TypeError as e:\n    if 'fwd output type' in str(e):\n        # align fwd return avals with primal outputs and retry\n        ...","preventionTips":["Write fwd rules that literally call the primal computation for the primal part","Use jax.eval_shape to compare avals in unit tests","Beware x64 mode flips changing dtypes"],"tags":["jax","custom-vjp","autodiff","type-mismatch"],"backgroundTag":"jax-custom-vjp-fwd-type-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}