jax-ml/jax · error · TypeError

Can't use ``defjvps`` with ``nondiff_argnums``.

Error message

Can't use ``defjvps`` with ``nondiff_argnums``.

What it means

The defjvps convenience method (per-argument JVP definitions) only works when no nondiff_argnums were specified on the @jax.custom_jvp decorator, because defjvps builds a rule that differentiates with respect to every positional argument. Combining the two is a TypeError.

Source

Thrown at jax/_src/hijax.py:1183

class custom_jvp3:
  jvp_fun: Callable | None = None
  symz: bool = False

  def __init__(self, f, nondiff_argnums=(), nondiff_argnames=()):
    self.static_argnums = _set_up_nondiff(f, nondiff_argnums, nondiff_argnames)
    update_wrapper(self, f)
    self.f = f

  def defjvp(self, jvp, symbolic_zeros=False):
    self.jvp_fun = jvp
    self.symz = symbolic_zeros
    return jvp

  def defjvps(self, *jvps):
    if self.static_argnums:
      raise TypeError("Can't use ``defjvps`` with ``nondiff_argnums``.")
    def jvp(primals, tangents):
      primal_out = self(*primals)
      zeros = tree_map(ad_util.p2tz, primal_out)
      all_tangents_out = [j(t, primal_out, *primals) if j else zeros
                          for t, j in zip(tangents, jvps)]
      sum_tangents = lambda _, x, *xs: reduce(ad.add_tangents, xs, x)
      tangent_out = tree_map(sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out
    self.defjvp(jvp)

  def __call__(self, *args, **kwargs):
    if not self.jvp_fun:
      msg = (f"No JVP defined for custom_jvp function {self.f.__name__} "
             "using defjvp.")
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.f, args, kwargs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use defjvp with a single full JVP function instead of defjvps
2. Remove nondiff_argnums from the decorator if per-argument definition is desired
3. Move static parameters outside the function (close over them or use functools.partial)

Example fix

// before
@jax.custom_jvp(nondiff_argnums=(1,))
def f(x, n): ...
f.defjvps(lambda t, p, x: 2*t)   # TypeError
// after
@jax.custom_jvp
def f(x, n): ...
@f.defjvp
def f_jvp(primals, tangents):
  (x, n), (xd, nd) = primals, tangents
  return f(x, n), 2 * xd

if getattr(f, 'static_argnums', None):
    raise TypeError('use defjvp, not defjvps, with nondiff_argnums')
f.defjvps(*jvps)

When it happens

Trigger: Calling f.defjvps(jvp1, jvp2, ...) on a function decorated as @jax.custom_jvp(nondiff_argnums=(1,)).

Common situations: Porting code between the defjvp and defjvps styles; adding nondiff_argnums later for static config without removing defjvps usage.

Related errors

• Can't use ``defjvps`` with ``nondiff_argnums``.
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• check_error takes an Error as argument, got type {type(error

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