{"record":{"id":"fb4f52632397b42b","repo":"jax-ml/jax","slug":"custom-gradient-function-used-with-with-logs-true-fb4f52","errorCode":null,"errorMessage":"custom_gradient function used with with_logs=True must return a VJP function whose second output is None or a dict of backward-pass log entries, but got {type(logs).__name__}.","messagePattern":"custom_gradient function used with with_logs=True must return a VJP function whose second output is None or a dict of backward-pass log entries, but got (.+?)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":1311,"sourceCode":"\n  if with_logs:\n    wrapped_fun.defvjp_with_logs(fwd, bwd)\n  else:\n    wrapped_fun.defvjp(fwd, bwd)\n  return wrapped_fun\n\ndef _custom_gradient_logs_rule(rule):\n  @wraps(rule)\n  def rule_with_logs(*cts):\n    out = rule(*cts)\n    if not (isinstance(out, (list, tuple)) and len(out) == 2):\n      raise TypeError(\n          \"custom_gradient function used with with_logs=True must return a \"\n          \"VJP function producing a pair (in_cts, logs), but the VJP function \"\n          f\"returned {out}.\")\n    in_cts, logs = out\n    if logs is not None and type(logs) is not dict:\n      raise TypeError(\n          \"custom_gradient function used with with_logs=True must return a \"\n          \"VJP function whose second output is None or a dict of \"\n          f\"backward-pass log entries, but got {type(logs).__name__}.\")\n    return in_cts, logs\n  return rule_with_logs\n\n@register_pytree_node_class\nclass Residuals:\n  def __init__(self, jaxpr, in_tree, out_tree, consts):\n    self.jaxpr = jaxpr\n    self.in_tree = in_tree\n    self.out_tree = out_tree\n    self.consts = consts\n  def __iter__(self):\n    return iter((self.jaxpr, self.in_tree, self.out_tree, self.consts))\n  def tree_flatten(self):\n    return self.consts, (self.jaxpr, self.in_tree, self.out_tree)\n  @classmethod","sourceCodeStart":1293,"sourceCodeEnd":1329,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L1293-L1329","documentation":"jax.custom_gradient with with_logs=True requires the VJP function to return a pair (in_cts, logs) where logs is None or a dict. This TypeError is raised when the second element of the returned pair is neither None nor a dict.","triggerScenarios":"Calling a function decorated with @custom_gradient(..., with_logs=True) and taking gradients, where the returned bwd function returns e.g. a tuple, list, array, or string instead of a dict or None as the second output.","commonSituations":"Porting a custom_gradient from the logs-free API and forgetting to update the VJP return shape; returning (cts, None, extra) triples; returning log arrays instead of dicts.","solutions":["Make the VJP function return exactly (in_cts, logs_dict) where logs_dict is a dict of backward-pass entries","Return (in_cts, None) if no backward logs are needed","If logs are not required, drop with_logs=True from the decorator"],"exampleFix":"# before\n@custom_gradient(with_logs=True)\ndef f(x):\n  def vjp(g):\n    return g * 2, [\"log\"]  # list is invalid\n  return x * 2, vjp\n\n# after\n@custom_gradient(with_logs=True)\ndef f(x):\n  def vjp(g):\n    return g * 2, {\"grad_norm\": jnp.linalg.norm(g)}\n  return x * 2, vjp","handlingStrategy":"validation","validationCode":"def check_vjp_shape(bwd):\n    cts, logs = bwd(jnp.ones(n))\n    assert logs is None or type(logs) is dict, type(logs)\n    return bwd","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Write a smoke test that calls the returned vjp with dummy cotangents and asserts the second output is dict or None"],"tags":["jax","custom-gradient","typeerror","vjp"],"backgroundTag":"custom-gradient-return-shape","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}