Redefines value_and_pullback_function to return a value and a pullback function

Project.toml

@@ -1,7 +1,7 @@
 name = "AbstractDifferentiation"
 uuid = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
 authors = ["Mohamed Tarek <mohamed82008@gmail.com> and contributors"]
-version = "0.3.0"
+version = "0.4.0-DEV"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

README.md

@@ -58,7 +58,7 @@ This operation goes by a few names. Refer to the [ChainRules documentation](http
 
 The following functions can be used to request the pullback operator/function with or without the function value. In order to request the pullback function `pb_f` of a function `f` at the inputs `xs`, you can use either of:
 - `pb_f = AD.pullback_function(ab::AD.AbstractBackend, f, xs...)`: returns the pullback function `pb_f` of the function `f` at the inputs `xs`. `pb_f` is a function that accepts the co-tangents `vs` as input which is a tuple of length equal to the number of outputs of `f`. If `f` has a single output, `pb_f` can also accept a single input instead of a 1-tuple.
-- `value_and_pb_f = AD.value_and_pullback_function(ab::AD.AbstractBackend, f, xs...)`: returns a function `value_and_pb_f` which accepts the co-tangent `vs` as input which is a tuple of length equal to the number of outputs of `f`. If `f` has a single output, `value_and_pb_f` can accept a single input instead of a 1-tuple. `value_and_pb_f` returns a 2-tuple, namely the value `f(xs...)` and output of the pullback operator.
+- `value_and_pb_f = AD.value_and_pullback_function(ab::AD.AbstractBackend, f, xs...)`: computes the function value `v = f(xs...)` and returns a 2-tuple containing the value `v` and a function `pb_f` that accepts the co-tangent `vs` as input, which is a tuple of length equal to the number of outputs of `f`. If `f` has a single output, `pb_f` can accept a single input instead of a 1-tuple.
 
 ### Lazy operators
 

src/AbstractDifferentiation.jl

@@ -249,31 +249,9 @@ function value_and_pullback_function(
     f,
     xs...,
 )
-    return (ws) -> begin
-        local value
-        primalcalled = false
-        if ab isa AbstractFiniteDifference
-            value = primal_value(ab, nothing, f, xs)
-            primalcalled = true
-        end
-        if ws === nothing
-            vs = f(xs...)
-            if !primalcalled
-                value = primal_value(lowest(ab), vs, f, xs)
-                primalcalled = true
-            end
-            return value, nothing
-        end
-        pb = pullback_function(lowest(ab), (_xs...,) -> begin
-            vs = f(_xs...)
-            if !primalcalled
-                value = primal_value(lowest(ab), vs, f, xs)
-                primalcalled = true
-            end
-            return vs
-        end, xs...)(ws)
-        return value, pb
-    end
+    # TODO: use value_and_pullback_function as a primitive to avoid double execution of primal
+    # https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/34
+    return f(xs...), pullback_function(lowest(ab), f, xs...)
 end
 
 struct LazyDerivative{B, F, X}
@@ -564,18 +542,17 @@ function define_pullback_function_and_friends(fdef)
     funcs = quote
         $(ExprTools.combinedef(fdef))
         function AbstractDifferentiation.jacobian($(args...),)
-            value_and_pbf = AbstractDifferentiation.value_and_pullback_function($(args...),)
-            value, _ = value_and_pbf(nothing)
+            value, pbf = AbstractDifferentiation.value_and_pullback_function($(args...),)
             identity_like = AbstractDifferentiation.identity_matrix_like(value)
             if eltype(identity_like) <: Tuple{Vararg{AbstractMatrix}}
                 return map(identity_like) do identity_like_i
                     if VERSION < v"1.3"
                         return reduce(vcat, map(AbstractDifferentiation._eachcol.(identity_like_i)...) do (cols...)
-                            value_and_pbf(cols)[2]'
+                            pbf(cols)'
                         end)
                     else
                     return mapreduce(vcat, AbstractDifferentiation._eachcol.(identity_like_i)...) do (cols...)
-                        value_and_pbf(cols)[2]'
+                        pbf(cols)'
                         end
                     end
                 end
@@ -584,10 +561,10 @@ function define_pullback_function_and_friends(fdef)
                 # value is a (grad,). Then, identity_like is a (matrix,).
                 # cols loops over columns of the matrix  
                 return vcat.(mapslices(identity_like[1], dims=1) do cols
-                    adjoint.(value_and_pbf((cols,))[2])
+                    adjoint.(pbf((cols,)))
                 end ...)
             else
-                return adjoint.(value_and_pbf(identity_like)[2])
+                return adjoint.(pbf(identity_like))
             end
         end
     end

test/test_utils.jl

@@ -192,16 +192,19 @@ function test_j′vp(backend; multiple_inputs=true, rng=Random.GLOBAL_RNG)
     w = rand(rng, length(fjac(xvec, yvec)))
     if multiple_inputs
         pb1 = AD.pullback_function(backend, fjac, xvec, yvec)(w)
-        valvec, pb2 = AD.value_and_pullback_function(backend, fjac, xvec, yvec)(w)
+        valvec, pbf2 = AD.value_and_pullback_function(backend, fjac, xvec, yvec)
+        pb2 = pbf2(w)
         @test valvec == fjac(xvec, yvec)
         @test norm.(pb2 .- pb1) == (0, 0)
         @test minimum(isapprox.(pb1, (vJxp(xvec,yvec,w), vJyp(xvec,yvec,w)), atol=1e-10))
         @test xvec == xvec2
         @test yvec == yvec2
     end
 
-    valvec1, pb1 = AD.value_and_pullback_function(backend, x -> fjac(x, yvec), xvec)(w)
-    valvec2, pb2 = AD.value_and_pullback_function(backend, y -> fjac(xvec, y), yvec)(w)
+    valvec1, pbf1 = AD.value_and_pullback_function(backend, x -> fjac(x, yvec), xvec)
+    pb1 = pbf1(w)
+    valvec2, pbf2 = AD.value_and_pullback_function(backend, y -> fjac(xvec, y), yvec)
+    pb2 = pbf2(w)
     @test valvec1 == fjac(xvec, yvec)
     @test valvec2 == fjac(xvec, yvec)
     @test minimum(isapprox.((pb1[1],pb2[1]), (vJxp(xvec,yvec,w), vJyp(xvec,yvec,w)), atol=1e-10))