{"record":{"id":"63e6b1e46f664735","repo":"jax-ml/jax","slug":"result-the-bwd-rule-attached-to-primal-sourceinf","errorCode":null,"errorMessage":"{result}the bwd rule attached to {primal_sourceinfo} produced an output of type {ct_aval.str_short()} which doesn't match expected type {expected.str_short()}","messagePattern":"(.+?)the bwd rule attached to (.+?) produced an output of type (.+?) which doesn't match expected type (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":906,"sourceCode":"def _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\"\n        f\" type {expected.str_short()}\")\n\ndef _replace_none(primal_in_aval, maybe_ct):\n  if maybe_ct is None:\n    return ad_util.Zero(primal_in_aval.to_ct_aval())\n  else:\n    return maybe_ct\n\nclass custom_vjp3:\n  fwd: Callable | None = None\n  bwd: Callable | None = None\n  symz: bool = False\n  opt_remat: bool = False\n  with_logs: bool = False\n\n  def __init__(self, f, nondiff_argnums=(), nondiff_argnames=()):","sourceCodeStart":888,"sourceCodeEnd":924,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L888-L924","documentation":"The backward (VJP) rule attached to a @jax.custom_vjp function via defvjp returned cotangents whose avals don't match the expected cotangent avals of the primal inputs. JAX validates each bwd output type against primal_aval.to_ct_aval() (skippable by setting JAX_DISABLE_BWD_CHECKS or the config disable_bwd_checks).","triggerScenarios":"A bwd function returning arrays with wrong shape, dtype (e.g. float64 vs float32), or wrong pytree structure for the input cotangents; returning zeros of the wrong shape; integer float0 handling mistakes.","commonSituations":"Hand-written bwd rules that return jnp.zeros_like(residual) instead of zeros matching the input; enabling/disabling 64-bit mode after writing the rule; returning None where a cotangent is expected or vice versa.","solutions":["Make bwd return cotangents matching each diff-argnum input's shape and dtype (use jnp.zeros(x.shape, x.dtype) pattern)","Cast bwd outputs with jax.lax.convert_element_type to the input dtype","As a last-resort diagnostic escape hatch, set jax.config.update('jax_disable_bwd_checks', True) (does not fix the underlying wrong gradients)","Verify pytree structure of bwd's return matches the diff-argnums prefix of f's inputs"],"exampleFix":"// before\ndef f_bwd(res, ct):\n  return (jnp.zeros_like(ct.astype(jnp.float64)),)  # wrong dtype\n// after\ndef f_bwd(res, ct):\n  return (ct * res,)  # same shape/dtype as the input x","handlingStrategy":"validation","validationCode":"# check bwd cotangent avals against inputs\nexpected = jax.eval_shape(lambda *args: args[:-0] if False else args, *diff_args)\nct_out = jax.eval_shape(lambda res, ct: f_bwd(res, ct), residuals, dummy_ct)\n# structures/dtypes must match diff args","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept ValueError as e:\n    if 'bwd rule' in str(e) and \"doesn't match expected type\" in str(e):\n        # fix bwd return shape/dtype to match inputs\n        ...","preventionTips":["Mirror each bwd output to its input's shape/dtype with zeros_like patterns","Add eval_shape-based tests for fwd/bwd rules","Avoid casting inside bwd unless you cast back"],"tags":["jax","custom-vjp","autodiff","type-mismatch"],"backgroundTag":"jax-custom-vjp-bwd-type-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}