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:
{disagreements}

What it means

Multi-output variant of the tangent check: JAX zips primal, expected-tangent, and actual-tangent avals and lists every disagreement where the rule's tangent does not typematch the expected tangent type of the primal; the combined message enumerates each mismatched output.

Source

Thrown at jax/_src/hijax.py:1162

  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:
      disagreements = "\n".join(
          f"  primal {av_p.str_short()} with tangent {av_t.str_short()}, "
          f"expecting tangent {av_et}"
          for av_p, av_et, av_t in zip(primal_avals_out,
                                       expected_tangent_avals_out,
                                       tangent_avals_out)
          if not core.typematch(av_et, av_t))
      raise TypeError(
          "Custom JVP rule must produce primal and tangent outputs with "
          f"corresponding shapes and dtypes, but got:\n{disagreements}")


class custom_jvp3:
  jvp_fun: Callable | None = None
  symz: bool = False

  def __init__(self, f, nondiff_argnums=(), nondiff_argnames=()):
    self.static_argnums = _set_up_nondiff(f, nondiff_argnums, nondiff_argnames)
    update_wrapper(self, f)
    self.f = f

  def defjvp(self, jvp, symbolic_zeros=False):
    self.jvp_fun = jvp
    self.symz = symbolic_zeros
    return jvp

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Fix each listed mismatch; the error names primal, its tangent, and the expected tangent per output
2. Keep tangent construction symmetric with primal construction (build both in the same tree_map)
3. For non-differentiable outputs return the appropriate zero (e.g. ad.InstantiatedZero / None per API)

Example fix

// before
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return (x*2, x+1), (xd, jnp.ones(3))  # second tangent wrong shape
// after
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return (x*2, x+1), (2*xd, xd)

outs, tans = f_jvp(primals, tangents)
jax.tree_util.tree_map(lambda p, t: (p.shape == t.shape), outs, tans)

When it happens

Trigger: A custom_jvp function returning several outputs where one or more tangents have wrong shape/dtype — e.g. swapping the order of tangents relative to primals, or returning zeros with wrong dtype for one output.

Common situations: Tuple outputs whose tangent tuple elements are permuted or truncated; mixed float/int outputs where int outputs need float0/zero tangents; large refactors of multi-output functions.

Related errors

• Custom JVP rule must produce primal and tangent outputs with
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Can't use ``defjvps`` with ``nondiff_argnums``.
• No JVP defined for custom_jvp function {primal_name} using d
• The input arguments to the custom_jvp-decorated function {pr

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