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

Single-output case of the tangent-check: the tangent returned by a custom_jvp rule must have the shape/dtype equal to the tangent type of the primal output (e.g. float32 tangent for a float32 primal). The message reports expected vs got.

Source

Thrown at jax/_src/custom_derivatives.py:370

           "instead the JVP rule output's first element had shapes/dtypes of:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_jvp-decorated function {primal_name} had output "
           "shapes/dtypes of:\n"
           f"""    {str(ty_tree_).replace("'", "")}""")
      raise TypeError(m)
  primal_avals_out = [core.typeof(x).strip_weak_type() for x in primals_out]
  expected_tangent_avals_out = [
    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')

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return a tangent with the same shape and dtype as the primal output (use jnp.zeros_like(primal_out) as the base)
2. Cast/correct explicitly and re-run jax.grad

Example fix

# before
return y, 0
# after
return y, jnp.zeros_like(y)

expected = jax.eval_shape(lambda x: jax.zeros(x.shape, x.dtype), primal_out)
# ensure tangent.shape == primal_out.shape and tangent.dtype == primal_out.dtype before returning

When it happens

Trigger: Returning an integer-dtype tangent, a differently shaped tangent (extra/missing batch dim), or a Python scalar where an array is expected, for a single-output custom_jvp function.

Common situations: Rules that compute gradients in a different precision; tangent accidentally shaped like the input instead of the output; zeros created with the wrong dtype.

Related errors

• Custom JVP rule {jvp_name} for function {primal_name} must p
• Custom JVP rule must produce primal and tangent outputs with
• {name} was requested to map its argument along axis {axis},
• {''.join(msg)[:-2]}
• primal and tangent arguments to jax.jvp do not match; dtypes

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