jax-ml/jax · error · TypeError

primal and tangent arguments to jax.jvp must be tuples or…

Error message

primal and tangent arguments to jax.jvp must be tuples or lists; found {type(primals).__name__} and {type(tangents).__name__}.

What it means

jax.jvp requires primals and tangents as explicit sequences because tangents are matched position-by-position to primals. This fires before flattening when either argument is a scalar, dict, or other non-list/tuple container.

Solutions

1. Wrap primals and tangents in tuples/lists: jax.jvp(f, (x,), (t,))
2. If splatting, use jax.jvp(f, *args, **kwargs)-safe patterns or build tuples explicitly
3. Check both containers, since either being wrong triggers the error

Example fix

# before
jax.jvp(f, x, jnp.ones_like(x))
# after
jax.jvp(f, (x,), (jnp.ones_like(x),))

assert isinstance(primals, (tuple, list)) and isinstance(tangents, (tuple, list)), 'jvp needs tuple/list primals and tangents'

def jvp_args_ok(p, t): return isinstance(p, (tuple, list)) and isinstance(t, (tuple, list))

When it happens

Trigger: jax.jvp(f, x, t) where x is a single jnp array rather than (x,); passing generators, dicts, or numpy arrays as the primal/tangent containers.

Common situations: Forgetting to wrap a single argument in a tuple; refactoring from grad (single-argument style) to jvp; passing *args splats incorrectly.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• Called with a float0 array. float0s do not support any…
• Called with a float0 at position . float0s do not support…
• Can't use ``defjvps`` with ``nondiff_argnums``.
• interpret with dynamic grid bounds unsupported
• jvp called with different primal and tangent shapes;Got…

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

Appendix: source

Thrown at jax/_src/api.py:1448

    ``tangents_out`` value has the same Python tree structure and shapes as
    ``primals_out``. If ``has_aux`` is ``True``, returns a
    ``(primals_out, tangents_out, aux)`` tuple where ``aux``
    is the auxiliary data returned by ``fun``.

  For example:

  >>> import jax
  >>>
  >>> primals, tangents = jax.jvp(jax.numpy.sin, (0.1,), (0.2,))
  >>> print(primals)
  0.09983342
  >>> print(tangents)
  0.19900084
  """
  check_callable(fun)
  if (not isinstance(primals, (tuple, list)) or
      not isinstance(tangents, (tuple, list))):
    raise TypeError("primal and tangent arguments to jax.jvp must be tuples or lists; "
                    f"found {type(primals).__name__} and {type(tangents).__name__}.")
  return _jvp(fun, primals, tangents, has_aux=has_aux)

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 "

View on GitHub (pinned to 1e1c6a8fc0)