jax-ml/jax · error · TypeError

Custom JVP rule must produce primal and tangent outputs with

Error message

Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:\n{}

What it means

Multi-output version of the tangent shape/dtype check: for at least one output leaf, the returned tangent does not match the tangent type of the corresponding primal output; the message lists each disagreement.

Source

Thrown at jax/_src/custom_derivatives.py:378

    core.typeof(x).strip_weak_type().to_tangent_aval()
    for x in primals_out]
  tangent_avals_out = [core.typeof(t).strip_weak_type()
                       if type(t) is not SymbolicZero else t.aval.strip_weak_type()
                       for t in tangents_out]
  if not all(map(core.typematch, expected_tangent_avals_out, tangent_avals_out)):
    if len(expected_tangent_avals_out) == 1:
      (av_p,), (av_et,), (av_t,) = primal_avals_out, expected_tangent_avals_out, tangent_avals_out
      msg = ("Custom JVP rule must produce primal and tangent outputs with "
             "corresponding shapes and dtypes. Expected {} (tangent type of {}) but got {}.")
      raise TypeError(msg.format(av_et.str_short(), av_p.str_short(), av_t.str_short()))
    else:
      msg = ("Custom JVP rule must produce primal and tangent outputs with "
             "corresponding shapes and dtypes, but got:\n{}")
      disagreements = (
          f"  primal {av_p.str_short()} with tangent {av_t.str_short()}, expecting tangent {av_et}"
          for av_p, av_et, av_t in zip(primal_avals_out, expected_tangent_avals_out, tangent_avals_out)
          if av_et != av_t)
      raise TypeError(msg.format('\n'.join(disagreements)))
  store.store((out_tree, primal_avals, ()))
  return primals_out + tangents_out

class CustomJVPCallPrimitive(core.Primitive):
  multiple_results = True
  skip_canonicalization = True

  def bind_with_trace(self, trace, args, avals, params, /):
    params = dict(params)
    fun, jvp = params.pop('subfuns')
    return trace.process_custom_jvp_call(self, fun, jvp, args, **params)

  def impl(self, fun, _, *args):
    raise NotImplementedError

  def is_high(self, *_, call_jaxpr, **__):
    return call_jaxpr.is_high

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Build tangents leaf-by-leaf from the primal outputs (tree_map over the primal pytree)
2. For non-differentiable outputs, ensure the tangent matches their tangent type (often zeros of the corresponding dtype)
3. Run jax.eval_shape on the rule to inspect every output leaf

Example fix

# before
return (y, aux), (dy, d_aux_float)
# after
return (y, aux), (dy, zeros_tangent_of(aux))

import jax
ok = all(a.shape == t.shape and a.dtype == t.dtype
         for a, t in zip(jax.tree_util.tree_leaves(primal_out),
                         jax.tree_util.tree_leaves(tangent_out)))

When it happens

Trigger: A custom_jvp rule over a multi-output function where one or more tangents have wrong shape/dtype — e.g. tangent pytree built per-input rather than per-output.

Common situations: Functions returning (value, aux) where aux is integer/non-differentiable and the rule returns a float tangent for it; partially updated rules after adding a new output.

Related errors

• Custom JVP rule {jvp_name} for function {primal_name} must p
• Custom JVP rule {jvp_name} for function {primal_name} must p
• primal and tangent arguments to jax.jvp must have the same t
• structure of the differentiated function {jaxpr.debug_info.f
• Pure callbacks do not support JVP. Please use `jax.custom_jv

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