jax-ml/jax · error · ValueError

jvp called with different primal and tangent shapes;Got prim

Error message

jvp called with different primal and tangent shapes;Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}

What it means

Within jax.jvp, each primal leaf and its corresponding tangent leaf must have identical shapes. This error reports both shapes when they differ (dtype/tree checks already passed).

Source

Thrown at jax/_src/api.py:1469

def _jvp(fun: Callable, primals, tangents, has_aux=False):
  ps_ft = ft.flatten(primals)
  ts_ft = ft.flatten(tangents)
  if ps_ft.tree != ts_ft.tree:
    raise TypeError("primal and tangent arguments to jax.jvp must have the same tree "
                    f"structure; primals have tree structure {ps_ft.tree} whereas tangents have "
                    f"tree structure {ts_ft.tree}.")
  for p, t in zip(ps_ft, ts_ft):
    if not isinstance(core.typeof(p), ShapedArray): continue
    if core.primal_dtype_to_tangent_dtype(_dtype(p)) != _dtype(t):
      raise TypeError("primal and tangent arguments to jax.jvp do not match; "
                      "dtypes must be equal, or in case of int/bool primal dtype "
                      "the tangent dtype must be float0."
                      f"Got primal dtype {_dtype(p)} and so expected tangent dtype "
                      f"{core.primal_dtype_to_tangent_dtype(_dtype(p))}, but got "
                      f"tangent dtype {_dtype(t)} instead.")
    if np.shape(p) != np.shape(t):
      raise ValueError("jvp called with different primal and tangent shapes;"
                       f"Got primal shape {np.shape(p)} and tangent shape as {np.shape(t)}")

  out_primals, out_tangents, *aux = ad.jvp(fun, ps_ft, ts_ft, has_aux=has_aux)
  return out_primals.unflatten(), out_tangents.unflatten(), *aux

@overload
def linearize(fun: Callable, *primals, has_aux: Literal[False] = False,
              in_nzs: Any = None) -> tuple[Any, Callable]:
  ...

@overload
def linearize(fun: Callable, *primals, has_aux: Literal[True],
              in_nzs: Any = None) -> tuple[Any, Callable, Any]:
  ...

@partial(api_boundary, repro_api_name="jax.linearize")
def linearize(fun: Callable, *primals, has_aux: bool = False,
              in_nzs: Any = None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Generate tangents via jax.tree.map(jnp.zeros_like, primals) so shapes match by construction
2. Fix the tangent construction to use the corresponding primal's shape
3. Add a shape equality assert before the jvp call

Example fix

# before
jax.jvp(f, (w,), (jnp.zeros_like(w.T),))
# after
jax.jvp(f, (w,), (jnp.zeros_like(w),))

for p, t in zip(tree_leaves(primals), tree_leaves(tangents)):
    if hasattr(p, 'shape'):
        assert np.shape(p) == np.shape(t), f'shape mismatch {np.shape(p)} vs {np.shape(t)}'

def shapes_match(p, t): return not hasattr(p, 'shape') or np.shape(p) == np.shape(t)

When it happens

Trigger: jax.jvp(f, (jnp.zeros((2,3)),), (jnp.zeros((3,2)),)); tangents built from a differently-shaped array (e.g. transposed or reshaped) than primals.

Common situations: Tangent constructed from a different variable than its primal after refactors; broadcasting assumptions where shapes happen to differ in one dimension; transposed weights in custom gradients.

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Custom JVP rule {jvp_name} for function {primal_name} must p

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