Repository: google/jax

Issue:
Possible issue with partial gradient application
I'm trying to define the derivative of a function at a point, for then compiling it with jit, but either I'm not doing it in the right way (with functools.partial) or there is a problem somewhere.

Define a simple function of x and y and its derivative with respect to x:
```python
from jax import grad

def f(x,y):
    return x+y

df = grad(f, argnums=0)
```

This works:
```python
from functools import partial

partial(f, x = 0.)(y = 1)
# 1.0
```

This doesn't:
```python
from functools import partial

partial(df, x = 0.)(y = 1)
# IndexError: tuple index out of range
```

why? 


Discussion:
Thanks for the question!

We should provide a better error message here, but essentially it's because we're trying to differentiate with respect to the first positional parameter (`argnums=0`) but we're not passing any positional arguments (both are passed as keywords).

So this works:

```python
partial(df, 0.)(y = 1)
```

You might think that we could identify keyword arguments with their positions, but unfortunately there's no robust way to do that in Python: it would require inspecting argument signatures, but signatures can be obscured by decorators and the like. The bottom line is that all `df` knows is that you're passing it two keyword arguments, yet its code expects there to be at least one positional argument when `argnums=0`, and `n+1` positional arguments when `argnums=n`.

We should improve this error message to something that explains at least `n+1` positional arguments must be passed by the caller when `argnums=n`, like we do for `jit` with `static_argnums`.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -388,6 +388,13 @@ def value_and_grad(fun, argnums=0, has_aux=False, holomorphic=False):
 
   @wraps(fun, docstr=docstr, argnums=argnums)
   def value_and_grad_f(*args, **kwargs):
+    max_argnum = argnums if type(argnums) is int else max(argnums)
+    if max_argnum >= len(args):
+      msg = ("differentiating with respect to argnums={} requires at least "
+             "{} positional arguments to be passed by the caller, but got only "
+             "{} positional arguments.")
+      raise TypeError(msg.format(argnums, max_argnum + 1, len(args)))
+
     f = lu.wrap_init(fun, kwargs)
     f_partial, dyn_args = _argnums_partial(f, argnums, args)
     if not has_aux:


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1253,5 +1253,17 @@ def test_repr(self):
     rep = repr(np.ones(()) + 1.)
     self.assertStartsWith(rep, 'DeviceArray')
 
+  def test_grad_without_enough_args_error_message(self):
+    # https://github.com/google/jax/issues/1696
+    def f(x, y): return x + y
+    df = api.grad(f, argnums=0)
+    self.assertRaisesRegexp(
+        TypeError,
+        "differentiating with respect to argnums=0 requires at least 1 "
+        "positional arguments to be passed by the caller, but got only 0 "
+        "positional arguments.",
+        lambda: partial(df, x=0.)(y=1.))
+
+
 if __name__ == '__main__':
   absltest.main()