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. Expected {} (tangent type of {}) but got {}.

What it means

For a custom_jvp function with a single output, the tangent output of the JVP rule must be the tangent-type of the primal output (same shape; float dtype if the primal is float; float0 if primal is non-differentiable). This TypeError reports the expected vs actual tangent aval.

Source

Thrown at jax/_src/hijax.py:1152

def _jvp_check_tangent_avals(self, out, out_tangent):
  strip = lambda a: a.strip_weak_type() if hasattr(a, 'strip_weak_type') else a
  out_flat = tree_leaves_checked(self.out_tree, out)
  tangents_flat = self.out_tree.flatten_up_to(out_tangent)
  primal_avals_out = [strip(typeof(x)) for x in out_flat]
  expected_tangent_avals_out = [a.to_tangent_aval() for a in primal_avals_out]
  tangent_avals_out = [
      strip(t.aval) if isinstance(t, (ad_util.Zero, ad_util.SymbolicZero))
      else strip(typeof(t)) for t in tangents_flat]
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match tangent shape/dtype to the primal: use jnp.ones_like / result * ones_like(primal) rather than new arrays with other dtypes
2. Wrap derivative constants as primal-dtype: 2.0 * xd or jnp.array(2.0, xd.dtype)
3. Test with jax.jvp(f, (x,), (t,)) for representative inputs

Example fix

// before
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return x * 2, 2        # scalar int tangent
// after
@f.defjvp
def f_jvp(p, t):
  (x,), (xd,) = p, t
  return x * 2, 2 * xd  # array tangent matching primal

av_p, av_t = jax.eval_shape(lambda p, t: (f_jvp(p, t)[0], f_jvp(p, t)[1]), (x,), (t,))
assert av_p.shape == av_t.shape and av_t.dtype == jnp.result_type(av_p.dtype, jnp.float32) or av_p.dtype == av_t.dtype

When it happens

Trigger: Returning a tangent with a different shape than the primal (e.g. forgetting to broadcast), an integer tangent for a float primal, or a scalar tangent for an array primal.

Common situations: Rules that return jnp.ones_like(primal) with a cast, integer-valued derivative expressions (e.g. from integer constants), tangent computed at wrong shapes when primal gets broadcast.

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/bbe452cbe5952ce9. Report an issue: GitHub.