Merge pull request #1419 from mcabbott/withgrad3

Allow `f` to return a Tuple in `withgradient(f, args...)`

Author: mcabbott

src/compiler/interface.jl

@@ -115,12 +115,34 @@ as a named tuple.
 julia> y, ∇ = withgradient(/, 1, 2)
 (val = 0.5, grad = (0.5, -0.25))
 
-julia> ∇ == gradient(/, 1, 2)  # explicit mode
+julia> ∇ == gradient(/, 1, 2)
 true
+```
+
+Allows you to capture auxillary outputs, in addition to the scalar
+used by `gradient`. To do this, `f` must return a Tuple or NamedTuple.
+Then it calculates `grad = gradient(first∘f, args...)
+but returns the whole `val = f(args...)`:
+
+```jldoctest; setup=:(using Zygote)
+julia> withgradient([1,2,4]) do x
+          z = 1 ./ x
+          sum(z), z  # here z is an auxillary output
+       end
+(val = (1.75, [1.0, 0.5, 0.25]), grad = ([-1.0, -0.25, -0.0625],))
+
+julia> withgradient(3.0, 4.0) do x, y
+          (div = x/y, mul = x*y)
+       end
+(val = (div = 0.75, mul = 12.0), grad = (0.25, -0.1875))
+```
+
+Also supports implicit mode:
 
+```jldoctest; setup=:(using Zygote)
 julia> w = [3.0];
 
-julia> res = withgradient(() -> sum(abs2, w), Params([w]))  # implicit mode
+julia> res = withgradient(() -> sum(abs2, w), Params([w]))
 (val = 9.0, grad = Grads(...))
 
 julia> res.grad[w]
@@ -130,7 +152,15 @@ julia> res.grad[w]
 """
 function withgradient(f, args...)
   y, back = pullback(f, args...)
-  grad = back(sensitivity(y))
+  grad = if y isa Tuple
+    dy = (sensitivity(first(y)), map(_ -> nothing, Base.tail(y))...)
+    back(dy)
+  elseif y isa NamedTuple
+    dy = (sensitivity(first(y)), map(_ -> nothing, Base.tail(y))...)
+    back(NamedTuple{propertynames(y), typeof(dy)}(dy))
+  else
+    back(sensitivity(y))
+  end
   results = isnothing(grad) ? map(_ -> nothing, args) : map(_project, args, grad)
   (val=y, grad=results)
 end

test/features.jl

@@ -866,3 +866,21 @@ end
   end
   @test gradient(f760, 3)[1] ≈ 123.93054835019153
 end
+
+@testset "withgradient" begin
+  @test withgradient([1,2,4]) do x
+    z = 1 ./ x
+    sum(z), z
+  end == (val = (1.75, [1.0, 0.5, 0.25]), grad = ([-1.0, -0.25, -0.0625],))
+
+  @test withgradient(3.0, 4.0) do x, y
+    (div = x/y, mul = x*y)
+  end == (val = (div = 0.75, mul = 12.0), grad = (0.25, -0.1875))
+
+  f3(x) = sum(sin, x), sum(cos, x), sum(tan, x)
+  g1 = gradient(first∘f3, [1,2,3.0])
+  y2, g2 = withgradient(first∘f3, [1,2,3.0])
+  y3, g3 = withgradient(f3, [1,2,3.0])
+  @test g1[1] ≈ g2[1] ≈ g3[1]
+end
+