jax-ml/jax · error · TypeError

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

Error message

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

What it means

jax.custom_jvp's defjvps convenience helper cannot be combined with nondiff_argnums, because defjvps assumes every positional argument is differentiable and gets a jvp entry. Supplying both is rejected with a TypeError.

Source

Thrown at jax/_src/custom_derivatives.py:243

      None.

    Examples:

      >>> @jax.custom_jvp
      ... def f(x, y):
      ...   return jnp.sin(x) * y
      ...
      >>> f.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y,
      ...           lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot)

      >>> x = jnp.float32(1.0)
      >>> y = jnp.float32(2.0)
      >>> with jnp.printoptions(precision=2):
      ...   print(jax.value_and_grad(f)(x, y))
      (Array(1.68, dtype=float32), Array(1.08, dtype=float32))
    """
    if self.nondiff_argnums:
      raise TypeError("Can't use ``defjvps`` with ``nondiff_argnums``.")

    def jvp(primals, tangents):
      primal_out = self(*primals)
      zeros = _zeros_like_pytree(primal_out)
      all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros
                          for t, jvp in zip(tangents, jvps)]
      tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out

    self.defjvp(jvp)

  @partial(traceback_util.api_boundary,
           repro_api_name="jax.custom_jvp.__call__")
  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:
    debug = debug_info("custom_jvp fun", self.fun, args, kwargs,
                       static_argnums=self.nondiff_argnums)
    primal_name = debug.func_name
    if not self.jvp:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch to the explicit defjvp form (f.defjvp) which handles nondiff args by receiving them as static values in the rule
2. Drop nondiff_argnums and pass the static values as arrays/traced values if semantically acceptable

Example fix

# before
@jax.custom_jvp(nondiff_argnums=(1,))
def f(x, n): ...
f.defjvps(lambda t, p, x, n: ...)
# after
@jax.custom_jvp
def f(x, n): ...
@f.defjvp
def f_jvp(primals, tangents):
    x, n = primals
    tx, _ = tangents
    return f(x, n), tx * dfdx(x, n)

# at decoration time choose one style:
# either @jax.custom_jvp (no nondiff_argnums) + defjvps, or nondiff_argnums + defjvp

When it happens

Trigger: Declaring @jax.custom_jvp(nondiff_argnums=(1,)) on a function and then calling f.defjvps(...) to register per-argument JVPs.

Common situations: Porting code that used static integer flags via nondiff_argnums while keeping a defjvps-style rule table; growing an existing custom_jvp function with extra static params.

Related errors

• primal and tangent arguments to jax.jvp must be tuples or li
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• No JVP defined for custom_jvp function {primal_name} using d
• Custom JVP rule {jvp_name} for function {primal_name} must p
• 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/852bbe3f49a1de71. Report an issue: GitHub.