jax-ml/jax · error · TypeError

custom_gradient function used with with_logs=True must retur

Error message

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__}.

What it means

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.

Source

Thrown at jax/_src/custom_derivatives.py:1311

  if with_logs:
    wrapped_fun.defvjp_with_logs(fwd, bwd)
  else:
    wrapped_fun.defvjp(fwd, bwd)
  return wrapped_fun

def _custom_gradient_logs_rule(rule):
  @wraps(rule)
  def rule_with_logs(*cts):
    out = rule(*cts)
    if not (isinstance(out, (list, tuple)) and len(out) == 2):
      raise TypeError(
          "custom_gradient function used with with_logs=True must return a "
          "VJP function producing a pair (in_cts, logs), but the VJP function "
          f"returned {out}.")
    in_cts, logs = out
    if logs is not None and type(logs) is not dict:
      raise TypeError(
          "custom_gradient function used with with_logs=True must return a "
          "VJP function whose second output is None or a dict of "
          f"backward-pass log entries, but got {type(logs).__name__}.")
    return in_cts, logs
  return rule_with_logs

@register_pytree_node_class
class Residuals:
  def __init__(self, jaxpr, in_tree, out_tree, consts):
    self.jaxpr = jaxpr
    self.in_tree = in_tree
    self.out_tree = out_tree
    self.consts = consts
  def __iter__(self):
    return iter((self.jaxpr, self.in_tree, self.out_tree, self.consts))
  def tree_flatten(self):
    return self.consts, (self.jaxpr, self.in_tree, self.out_tree)
  @classmethod

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the VJP function return exactly (in_cts, logs_dict) where logs_dict is a dict of backward-pass entries
2. Return (in_cts, None) if no backward logs are needed
3. If logs are not required, drop with_logs=True from the decorator

Example fix

# before
@custom_gradient(with_logs=True)
def f(x):
  def vjp(g):
    return g * 2, ["log"]  # list is invalid
  return x * 2, vjp

# after
@custom_gradient(with_logs=True)
def f(x):
  def vjp(g):
    return g * 2, {"grad_norm": jnp.linalg.norm(g)}
  return x * 2, vjp

def check_vjp_shape(bwd):
    cts, logs = bwd(jnp.ones(n))
    assert logs is None or type(logs) is dict, type(logs)
    return bwd

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• reduce_axes argument to vjp is deprecated
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass
• the gradient for {_vjp_arg_name(jaxpr, in_tree, idx)}, which
• {_vjp_arg_name(jaxpr, in_tree, idx)} is Ref-typed, so its gr
• But the tree structures differ:

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/fb4f52632397b42b. Report an issue: GitHub.