Make type of value_and_grad slightly more precise. (jax-ml#2704)

jacobjinkelly · Apr 21, 2020 · 5dac151 · 5dac151
1 parent 2603d5f
commit 5dac151
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/jax/api.py b/jax/api.py
@@ -380,7 +380,8 @@ def grad_f_aux(*args, **kwargs):
   return grad_f_aux if has_aux else grad_f
 
 def value_and_grad(fun: Callable, argnums: Union[int, Sequence[int]] = 0,
-                   has_aux: bool = False, holomorphic: bool = False) -> Callable:
+                   has_aux: bool = False, holomorphic: bool = False
+                   ) -> Callable[..., Tuple[Any, Any]]:
   """Create a function which evaluates both ``fun`` and the gradient of ``fun``.
 
   Args: