Repository: google/jax

Issue:
Inscrutable custom VJP error
### Description

Here's my MWE.  I left in the type annotations and the docstring in case they help with making sense of the code.
```python
from functools import partial
from typing import Any, Callable, Tuple, TypeVar, cast

import jax.numpy as jnp
from jax import custom_vjp, vjp
from jax.tree_util import tree_map

XT = TypeVar('XT', bound=Tuple[Any, ...])
Y = TypeVar('Y')

@partial(custom_vjp, nondiff_argnums=(0,))
def cotangent_combinator(f: Callable[..., Tuple[XT, Y]],
                         args_tuples: Tuple[Tuple[Any, ...], ...]) -> Tuple[XT, Y]:
    """
    Args:
        f: A function that accepts positional arguments and returns xs, y where xs is a tuple of
            length n.
        args_tuples: n copies of the same tuple of positional arguments accepted by f.
    Returns: The pair (xs, y).

    The purpose of the cotangent combinator is to take cotangents of each of the elements of x, and
    send them back through to each of the corresponding argument tuples.
    """
    return f(*args_tuples[0])


def _cotangent_combinator_fwd(f: Callable[..., Tuple[XT, Y]],
                              args_tuples: Tuple[Tuple[Any, ...], ...]
                              ) -> Tuple[Tuple[XT, Y], Callable[[XT, Y], Tuple[Any, ...]]]:
    return vjp(f, *args_tuples[0])


def _cotangent_combinator_bwd(f: Callable[..., Tuple[XT, Y]],
                              f_vjp: Callable[[XT, Y], Tuple[Any, ...]],
                              xy_bar: Tuple[XT, Y]
                              ) -> Tuple[Any, ...]:
    xs_bar, y_bar = xy_bar
    n = len(xs_bar)
    xs_zero = tuple(tree_map(jnp.zeros_like, x_bar)
                    for x_bar in xs_bar)
    all_args_bar = []
    for i, x_bar in enumerate(xs_bar):
        this_xs_bar = cast(XT, (xs_zero[:i]
                                + (x_bar,)
                                + xs_zero[i + 1:]))
        print_generic(this_xs_bar, y_bar)
        this_args_bar = f_vjp(this_xs_bar, y_bar)
        all_args_bar.append(this_args_bar)
    return tuple(all_args_bar)


cotangent_combinator.defvjp(_cotangent_combinator_fwd, _cotangent_combinator_bwd)


def f(x: Any) -> Tuple[Any, None]:
    return (x ** 2, x ** 2), None
o = jnp.ones(())

result, f_vjp = vjp(partial(cotangent_combinator, f), ((-1.0,), (-1.0,)))
result_bar = (2 * o, 3 * o), None
from tjax import print_generic
print_generic(result=result)
print_generic(result_bar=result_bar)
f_vjp(result_bar)
```
gives
```
TypeError: _vjp_pullback_wrapper() takes 5 positional arguments but 6 were given
```
I have no idea where this is coming from since I don't have anything like that in my code.  What am I doing wrong?

### What jax/jaxlib version are you using?

jax 0.3.23; jaxlib 0.3.22

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.10


Discussion:
A terrible error message indeed!

In the original repro (before edit, sorry that's just when I copied it) the issue was that we were calling

```python
f_vjp((2.0, 0.0), None)
```

when we actually needed

```python
f_vjp(((2.0, 0.0), None))
```

Notice the extra parentheses. I think `f_vjp` always takes a single argument (corresponding to how Python functions always return a single value, though that value might be a tuple or pytree or whatever).

Let me try the new repro... but if the issue is the same, then at least we know how to fix the code, though it's still a TODO to improve the error message here.
The callable returned by `jax.vjp` is a `Partial` over a function which takes 5 arguments (named `_vjp_pullback_wrapper`) where all but the last argument is supplied by the `Partial`. Because of the `Partial`, we don't error quite as early as you might expect (and apparently tracebacks are messed up, because we subclass `functools.partial` perhaps?).
I think in the new repro we have a repeat of the same issue on this line:

```python
        this_args_bar = f_vjp(this_xs_bar, y_bar)
```

where we instead want

```python
        this_args_bar = f_vjp((this_xs_bar, y_bar))
```
Yeah, you got it!  Thank you.  Please let me know if/when I should close this.
I'll send a PR with an improved error message momentarily. Let's let the PR close the issue. (Please review the PR if you feel like it!)

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2531,8 +2531,30 @@ def fun(*tangents):
 
   return apply_flat_fun(fun, io_tree, *py_args)
 
-def _vjp_pullback_wrapper(cotangent_dtypes, cotangent_shapes,
-                          io_tree, fun, py_args):
+def _vjp_pullback_wrapper(name, cotangent_dtypes, cotangent_shapes, io_tree,
+                          fun, *py_args_):
+  if len(py_args_) != 1:
+    msg = (f"The function returned by `jax.vjp` applied to {name} was called "
+           f"with {len(py_args_)} arguments, but functions returned by "
+           "`jax.vjp` must be called with a single argument corresponding to "
+           f"the single value returned by {name} (even if that returned "
+           "value is a tuple or other container).\n"
+           "\n"
+           "For example, if we have:\n"
+           "\n"
+           "  def f(x):\n"
+           "    return (x, x)\n"
+           "  _, f_vjp = jax.vjp(f, 1.0)\n"
+           "\n"
+           "the function `f` returns a single tuple as output, and so we call "
+           "`f_vjp` with a single tuple as its argument:\n"
+           "\n"
+           "  x_bar, = f_vjp((2.0, 2.0))\n"
+           "\n"
+           "If we instead call `f_vjp(2.0, 2.0)`, with the values 'splatted "
+           "out' as arguments rather than in a tuple, this error can arise.")
+    raise TypeError(msg)
+  py_args, = py_args_
   in_tree_expected, out_tree = io_tree
   args, in_tree = tree_flatten(py_args)
   if in_tree != in_tree_expected:
@@ -2637,9 +2659,8 @@ def _vjp(fun: lu.WrappedFun, *primals, has_aux=False, reduce_axes=()):
   ct_shapes = [np.shape(x) for x in out_primal]
   # Ensure that vjp_py is a PyTree so that we can pass it from the forward to the
   # backward pass in a custom VJP.
-  vjp_py = Partial(partial(_vjp_pullback_wrapper,
-                           ct_dtypes, ct_shapes,
-                           (out_tree, in_tree)),
+  vjp_py = Partial(partial(_vjp_pullback_wrapper, fun.__name__,
+                           ct_dtypes, ct_shapes, (out_tree, in_tree)),
                    out_vjp)
   if not has_aux:
     return out_primal_py, vjp_py


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -4095,6 +4095,14 @@ def h(x):
 
     b(8)  # don't crash
 
+  def test_vjp_multiple_arguments_error_message(self):
+    # https://github.com/google/jax/issues/13099
+    def foo(x):
+      return (x, x)
+    _, f_vjp = jax.vjp(foo, 1.0)
+    with self.assertRaisesRegex(TypeError, "applied to foo"):
+      f_vjp(1.0, 1.0)
+
 
 @jtu.with_config(jax_experimental_subjaxpr_lowering_cache=True)
 class SubcallTraceCacheTest(jtu.JaxTestCase):