jax-ml/jax · error · TypeError

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

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 to be Python tuples or lists (not scalars, arrays, dicts, or other containers). This check runs before flattening and rejects any other sequence type.

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)

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: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• primal and tangent arguments to jax.jvp must have the same t
• 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/7c7c1e1cda204f30. Report an issue: GitHub.