Repository: google/jax

Issue:
add a way to return side-information from grad
For example, a user may want to return a statistic or auxiliary data structure from the Python function to which `grad` is applied, not to be considered as part of the value of the mathematical function to be differentiated.


Discussion:
Following up from JAXers conversation: it's possible to do this by manually invoking `vjp` with more than one seed:
```python
import jax
def loss(x): return (x + x, x)
def mygrad(x):
  (l, extra), pullback = jax.vjp(loss, x)
  dx, = pullback((1.0, 0.0))
  return dx, extra
```
But as Dougal pointed out, even calling `jax.ad.vjp` with `jax.ad_util.zero` (rather than `0.0`) [doesn't avoid materializing](https://github.com/google/jax/blob/63608fb9715b33cfac64dc26548fcd7da35e97d3/jax/interpreters/ad.py#L79) the second seed (and performing unnecessary backward-pass computations).

So it would be nice to have an API that either accepts real symbolic zeros or allows selecting a single scalar output to receive a seed.
@jekbradbury as you probably know, that's [the approach taken in Autograd](https://github.com/HIPS/autograd/blob/master/autograd/differential_operators.py#L142..L147).

Note that the extra flop cost will, I think, not be too bad for most use cases, because often `extra` is something that's computed during computation of `l` anyway (or doesn't take much to compute from some intermediate value during the computation of `l`). This is the case in the trivial example you give. That means that the zeros will only be back-propagated through a small piece of graph before being merged into the gradient computation that's going on anyway. Does that make sense?

__EDIT:__ here's a couple of pictures to explain. Firstly the function itself:
```
inputs ---------+---------> outputs
                |
                |
                ↓
               aux
```
then the backpropagation looks like this
```
inputs <--------+---------- outputs
                ↑  
                |  < - zeros will only be propagated through this bit, 
                |  < - which is usually small
               aux
```

I still think there will probably be some cases where avoiding any extra FLOs would be beneficial.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -169,7 +169,7 @@ def computation_maker(*args, **kwargs):
 
   return computation_maker
 
-def grad(fun, argnums=0):
+def grad(fun, argnums=0, has_aux=False):
   """Creates a function which evaluates the gradient of `fun`.
 
   Args:
@@ -179,13 +179,17 @@ def grad(fun, argnums=0):
       arrays with shape `(1,)` etc.)
     argnums: Optional, integer or tuple of integers. Specifies which positional
       argument(s) to differentiate with respect to (default 0).
+    has_aux: Optional, bool. Indicates whether `fun` returns a pair where the
+     first element is considered the output of the mathematical function to be
+     differentiated and the second element is auxiliary data. Default False.
 
   Returns:
     A function with the same arguments as `fun`, that evaluates the gradient of
     `fun`. If `argnums` is an integer then the gradient has the same shape and
     type as the positional argument indicated by that integer. If argnums is a
     tuple of integers, the gradient is a tuple of values with the same shapes
-    and types as the corresponding arguments.
+    and types as the corresponding arguments. If `has_aux` is True then a pair
+    of (gradient, auxiliary_data) is returned.
 
   For example:
 
@@ -194,7 +198,7 @@ def grad(fun, argnums=0):
   array(0.961043, dtype=float32)
 
   """
-  value_and_grad_f = value_and_grad(fun, argnums)
+  value_and_grad_f = value_and_grad(fun, argnums, has_aux=has_aux)
 
   docstr = ("Gradient of {fun} with respect to positional argument(s) "
             "{argnums}. Takes the same arguments as {fun} but returns the "
@@ -203,12 +207,16 @@ def grad(fun, argnums=0):
 
   @wraps(fun, docstr=docstr, argnums=argnums)
   def grad_f(*args, **kwargs):
-    ans, g = value_and_grad_f(*args, **kwargs)
-    return g
+    if not has_aux:
+      _, g = value_and_grad_f(*args, **kwargs)
+      return g
+    else:
+      (_, aux), g = value_and_grad_f(*args, **kwargs)
+      return g, aux
 
   return grad_f
 
-def value_and_grad(fun, argnums=0):
+def value_and_grad(fun, argnums=0, has_aux=False):
   """Creates a function which evaluates both `fun` and the gradient of `fun`.
 
   Args:
@@ -218,6 +226,9 @@ def value_and_grad(fun, argnums=0):
       arrays with shape `(1,)` etc.)
     argnums: Optional, integer or tuple of integers. Specifies which positional
       argument(s) to differentiate with respect to (default 0).
+    has_aux: Optional, bool. Indicates whether `fun` returns a pair where the
+     first element is considered the output of the mathematical function to be
+     differentiated and the second element is auxiliary data. Default False.
 
   Returns:
     A function with the same arguments as `fun` that evaluates both `fun` and
@@ -238,11 +249,17 @@ def value_and_grad(fun, argnums=0):
   def value_and_grad_f(*args, **kwargs):
     f = lu.wrap_init(fun, kwargs)
     f_partial, dyn_args = _argnums_partial(f, argnums, args)
-    ans, vjp_py = vjp(f_partial, *dyn_args)
+    if not has_aux:
+      ans, vjp_py = vjp(f_partial, *dyn_args)
+    else:
+      ans, vjp_py, aux = vjp(f_partial, *dyn_args, has_aux=True)
     _check_scalar(ans)
     g = vjp_py(onp.ones((), onp.result_type(ans)))
     g = g[0] if isinstance(argnums, int) else g
-    return (ans, g)
+    if not has_aux:
+      return ans, g
+    else:
+      return (ans, aux), g
 
   return value_and_grad_f
 
@@ -529,7 +546,7 @@ def fun(*args):
 
   return apply_jaxtree_fun(fun, io_tree, *py_args)
 
-def vjp(fun, *primals):
+def vjp(fun, *primals, **kwargs):
   """Compute a (reverse-mode) vector-Jacobian product of `fun`.
 
   `grad` is implemented as a special case of `vjp`.
@@ -542,6 +559,9 @@ def vjp(fun, *primals):
       should be evaluated. The length of `primals` should be equal to the number
       of positional parameters to `fun`. Each primal value should be a tuple of
       arrays, scalar, or standard Python containers thereof.
+    has_aux: Optional, bool. Indicates whether `fun` returns a pair where the
+     first element is considered the output of the mathematical function to be
+     differentiated and the second element is auxiliary data. Default False.
 
   Returns:
     A `(primals_out, vjpfun)` pair, where `primals_out` is `fun(*primals)`.
@@ -556,20 +576,30 @@ def vjp(fun, *primals):
   >>> g((-0.7, 0.3))
   (array(-0.61430776, dtype=float32), array(-0.2524413, dtype=float32))
   """
+  has_aux = kwargs.pop('has_aux', False)
+  assert not kwargs
   if not isinstance(fun, lu.WrappedFun):
     fun = lu.wrap_init(fun)
   primals_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, primals))
   _check_args(primals_flat)
   jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(fun, in_trees)
-  out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)
+  if not has_aux:
+    out_primal, out_vjp = ad.vjp(jaxtree_fun, primals_flat)
+  else:
+    out_primal, out_vjp, aux = ad.vjp(jaxtree_fun, primals_flat, has_aux=True)
   out_tree = out_tree()
+  if has_aux:
+    out_tree, aux_tree = out_tree.children
   out_primal_py = build_tree(out_tree, out_primal)
   ct_in_trees = [out_tree]
   ct_out_tree = PyTreeDef(node_types[tuple], None, in_trees)
   def out_vjp_packed(cotangent_in):
     return out_vjp(cotangent_in)
   vjp_py = partial(apply_jaxtree_fun, out_vjp_packed, (ct_in_trees, ct_out_tree))
-  return out_primal_py, vjp_py
+  if not has_aux:
+    return out_primal_py, vjp_py
+  else:
+    return out_primal_py, vjp_py, build_tree(aux_tree, aux)
 
 
 def trace_to_jaxpr(traceable, py_pvals, **kwargs):
diff --git a/jax/interpreters/ad.py b/jax/interpreters/ad.py
--- a/jax/interpreters/ad.py
+++ b/jax/interpreters/ad.py
@@ -34,8 +34,12 @@
 def identity(x): return x
 
 
-def jvp(fun):
-  return jvpfun(jvp_subtrace(fun))
+def jvp(fun, has_aux=False):
+  if not has_aux:
+    return jvpfun(jvp_subtrace(fun))
+  else:
+    fun, aux = jvp_subtrace_aux(fun)
+    return jvpfun(fun), aux
 
 @transformation
 def jvpfun(primals, tangents):
@@ -56,13 +60,31 @@ def jvp_subtrace(master, primals, tangents):
   out_primal, out_tangent = out_tracer.primal, out_tracer.tangent
   yield (out_primal, out_tangent)
 
+@transformation_with_aux
+def jvp_subtrace_aux(master, primals, tangents):
+  trace = JVPTrace(master, core.cur_sublevel())
+  for x in list(primals) + list(tangents):
+    if isinstance(x, Tracer):
+      assert x.trace.level < trace.level
+  ans, aux = yield map(partial(JVPTracer, trace), primals, tangents)
+  out_tracer, aux_tracer = map(trace.full_raise, (ans, aux))
+  out_primal, out_tangent = out_tracer.primal, out_tracer.tangent
+  aux = aux_tracer.primal  # ignore aux tangent
+  yield (out_primal, out_tangent), aux
+
+
 @transformation
 def pack_output(*args):
   ans = yield args
   yield pack(ans)
 
-def linearize(traceable, *primals):
-  jvpfun = pack_output(jvp(traceable))
+def linearize(traceable, *primals, **kwargs):
+  has_aux = kwargs.pop('has_aux', False)
+  if not has_aux:
+    jvpfun = pack_output(jvp(traceable))
+  else:
+    jvpfun, aux = jvp(traceable, has_aux=True)
+    jvpfun = pack_output(jvpfun)
   tangent_avals = [get_aval(p).at_least_vspace() for p in primals]
   in_pvals = (pe.PartialVal((None, pack(primals))),
               pe.PartialVal((core.AbstractTuple(tangent_avals), core.unit)))
@@ -70,10 +92,16 @@ def linearize(traceable, *primals):
   pval_primal, pval_tangent = unpair_pval(out_pval)
   aval_primal, const_primal = pval_primal
   assert aval_primal is None
-  return const_primal, pval_tangent, jaxpr, consts
+  if not has_aux:
+    return const_primal, pval_tangent, jaxpr, consts
+  else:
+    return const_primal, pval_tangent, jaxpr, consts, aux()
 
-def vjp(traceable, primals):
-  out_primal, pval, jaxpr, consts = linearize(traceable, *primals)
+def vjp(traceable, primals, has_aux=False):
+  if not has_aux:
+    out_primal, pval, jaxpr, consts = linearize(traceable, *primals)
+  else:
+    out_primal, pval, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)
   def vjp_(ct):
     ct = ignore_consts(ct, pval)
     dummy_primal_and_ct = pack((core.unit, ct))
@@ -81,7 +109,10 @@ def vjp_(ct):
     _, arg_cts = backward_pass(jaxpr, consts, (), dummy_args, dummy_primal_and_ct)
     return instantiate_zeros(pack(primals), arg_cts[1])
 
-  return out_primal, vjp_
+  if not has_aux:
+    return out_primal, vjp_
+  else:
+    return out_primal, vjp_, aux
 
 def ignore_consts(ct, pval):
   aval, const = pval


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -341,6 +341,41 @@ def test_large_device_constant(self):
     ans = jit(lambda x: 2 * x)(np.ones(int(2e6)))  # doesn't crash
     self.assertAllClose(ans, 2., check_dtypes=False)
 
+  def test_grad_and_aux_basic(self):
+    g, aux = grad(lambda x: (x**3, [x**2]), has_aux=True)(3.)
+    self.assertEqual(g, grad(lambda x: x**3)(3.))
+    self.assertEqual(aux, [9.])
+
+  def test_grad_and_aux_nested(self):
+    def f(x):
+      g, aux = grad(lambda x: (x**3, [x**3]), has_aux=True)(x)
+      return aux[0]
+
+    f2 = lambda x: x**3
+
+    self.assertEqual(grad(f)(4.), grad(f2)(4.))
+    self.assertEqual(jit(grad(f))(4.), grad(f2)(4.))
+    self.assertEqual(jit(grad(jit(f)))(4.), grad(f2)(4.))
+
+    def f(x):
+      g, aux = grad(lambda x: (x**3, [x**3]), has_aux=True)(x)
+      return aux[0] * np.sin(x)
+
+    f2 = lambda x: x**3 * np.sin(x)
+
+    self.assertEqual(grad(f)(4.), grad(f2)(4.))
+    self.assertEqual(jit(grad(f))(4.), grad(f2)(4.))
+    self.assertEqual(jit(grad(jit(f)))(4.), grad(f2)(4.))
+
+  def test_grad_and_aux_constant(self):
+    g, aux = grad(lambda x: (x**3, [4.]), has_aux=True)(4.)
+    self.assertEqual(g, grad(lambda x: x**3)(4.))
+    self.assertEqual(aux, [4.])
+
+    g, aux = grad(lambda x: (x**3, [x**2, 4.]), has_aux=True)(4.)
+    self.assertEqual(g, grad(lambda x: x**3)(4.))
+    self.assertEqual(aux, [4.**2, 4.])
+
 
 if __name__ == '__main__':
   absltest.main()