jax-ml/jax · error · TypeError

Custom JVP rule {jvp_name} for function {primal_name} must p

Error message

Custom JVP rule {jvp_name} for function {primal_name} must produce a pair (list or tuple of length two) where the first element represents the primal output (equal in value to the output of the custom_jvp-decorated function {primal_name}, and in particular with leaves of the same shape/dtype), but instead the JVP rule output's first element had shapes/dtypes of:\n    {str(ty_tree ).replace("'", "")}\nwhile the custom_jvp-decorated function {primal_name} had output shapes/dtypes of:\n    {str(ty_tree_).replace("'", "")}

What it means

Stronger than the structure check: the leaves of a custom_jvp rule's primal output must have the same shape and dtype as the original function's outputs. This fires when structures match but an array's shape or dtype differs (e.g. broadcasting added a dim, or the rule computed in float64).

Source

Thrown at jax/_src/custom_derivatives.py:357

           "structure:\n"
           f"""    {str(ty_tree ).replace("'", "")}\n"""
           f"while the custom_jvp-decorated function {primal_name} had output "
           "container/pytree structure:\n"
           f"""    {str(ty_tree_).replace("'", "")}.""")
      raise TypeError(m)
    if not all(map(core.typematch, primal_avals, primal_avals_)):
      m = (f"Custom JVP rule {jvp_name} for function {primal_name} must "
           "produce a pair (list or tuple of length two) "
           "where the first element represents the primal output "
           "(equal in value to the output of the custom_jvp-decorated function "
           f"{primal_name}, "
           "and in particular with leaves of the same shape/dtype), but "
           "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}"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return arrays with exactly the primal's shape and dtype; cast with .astype(primal.dtype) if needed
2. Disable jax_enable_x64 or ensure constants are typed (jnp.float32(...)) to avoid promotion
3. Compare jax.eval_shape of the function and the rule's primal half

Example fix

# before
return y * 1.0, dy
# after
return y.astype(primal_dtype), dy

import jax
pspec = jax.eval_shape(f, *sample_args)
# in the rule: return jax.tree_util.tree_map(lambda a, b: b.astype(a.dtype), pspec, primal_out), tangent

When it happens

Trigger: A rule that returns y[None] instead of y, squeezes/expands dims, or promotes dtypes (e.g. multiplying by a Python float under x64 mode) relative to the primal output.

Common situations: x64 enabled causing silent dtype promotion; reshape/broadcast differences between the rule and the original implementation; stale rule after the primal's dtype contract changed.

Related errors

• Custom JVP rule must produce primal and tangent outputs with
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Can't use ``defjvps`` with ``nondiff_argnums``.

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