Add comments for residuals from f_bwd. (#4244)

docs/notebooks/Custom_derivative_rules_for_Python_code.ipynb

@@ -193,10 +193,11 @@
         "  return jnp.sin(x) * y\n",
         "\n",
         "def f_fwd(x, y):\n",
+        "# Returns primal output and residuals to be used in backward pass by f_bwd.\n",
         "  return f(x, y), (jnp.cos(x), jnp.sin(x), y)\n",
         "\n",
         "def f_bwd(res, g):\n",
-        "  cos_x, sin_x, y = res\n",
+        "  cos_x, sin_x, y = res # Gets residuals computed in f_fwd\n",
         "  return (cos_x * g * y, sin_x * g)\n",
         "\n",
         "f.defvjp(f_fwd, f_bwd)"
@@ -2820,4 +2821,4 @@
       ]
     }
   ]
-}
+}