change return type of grad/jacobian to match args (#51)

* change return type

* finish this

* revert jvp

* Update Project.toml

Author: Rogerluo

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDifferences"
 uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
-version = "0.8.0"
+version = "0.9.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/grad.jl

@@ -18,11 +18,11 @@ function grad(fdm, f, x::AbstractArray{T}) where T <: Number
             return f(tmp)
         end
     end
-    return dx
+    return (dx, )
 end
 
-grad(fdm, f, x::Real) = fdm(f, x)
-grad(fdm, f, x::Tuple) = grad(fdm, (xs...)->f(xs), x...)
+grad(fdm, f, x::Real) = (fdm(f, x), )
+grad(fdm, f, x::Tuple) = (grad(fdm, (xs...)->f(xs), x...), )
 
 function grad(fdm, f, d::Dict{K, V}) where {K, V}
     dd = Dict{K, V}()
@@ -32,19 +32,19 @@ function grad(fdm, f, d::Dict{K, V}) where {K, V}
             tmp[k] = x
             return f(tmp)
         end
-        dd[k] = grad(fdm, f′, v)
+        dd[k] = grad(fdm, f′, v)[1]
     end
-    return dd
+    return (dd, )
 end
 
 function grad(fdm, f, x)
     v, back = to_vec(x)
-    return back(grad(fdm, x->f(back(v)), v))
+    return (back(grad(fdm, x->f(back(v)), v)), )
 end
 
 function grad(fdm, f, xs...)
     return ntuple(length(xs)) do k
-        grad(fdm, x->f(replace_arg(x, xs, k)...), xs[k])
+        grad(fdm, x->f(replace_arg(x, xs, k)...), xs[k])[1]
     end
 end
 
@@ -57,17 +57,17 @@ Approximate the Jacobian of `f` at `x` using `fdm`. `f(x)` must be a length `len
 function jacobian(fdm, f, x::Union{T, AbstractArray{T}}; len::Int=length(f(x))) where {T <: Number}
     J = Matrix{float(T)}(undef, len, length(x))
     for d in 1:len
-        gs = grad(fdm, x->f(x)[d], x)
+        gs = grad(fdm, x->f(x)[d], x)[1]
         for k in 1:length(x)
             J[d, k] = gs[k]
         end
     end
-    return J
+    return (J, )
 end
 
 function jacobian(fdm, f, xs...; len::Int=length(f(xs...)))
     return ntuple(length(xs)) do k
-        jacobian(fdm, x->f(replace_arg(x, xs, k)...), xs[k]; len=len)
+        jacobian(fdm, x->f(replace_arg(x, xs, k)...), xs[k]; len=len)[1]
     end
 end
 
@@ -83,7 +83,7 @@ _jvp(fdm, f, x::Vector{<:Number}, ẋ::AV{<:Number}) = fdm(ε -> f(x .+ ε .* x
 
 Convenience function to compute `transpose(jacobian(f, x)) * ȳ`.
 """
-_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))) * ȳ
+_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))[1]) * ȳ
 
 """
     jvp(fdm, f, x, ẋ)
@@ -98,7 +98,7 @@ function jvp(fdm, f, (x, ẋ)::Tuple{Any, Any})
 end
 function jvp(fdm, f, xẋs::Tuple{Any, Any}...)
     x, ẋ = collect(zip(xẋs...))
-    return jvp(fdm, xs->f(xs...), (x, ẋ))
+    return jvp(fdm, xs->f(xs...)[1], (x, ẋ))
 end
 
 """
@@ -109,9 +109,9 @@ Compute an adjoint with any types of arguments for which [`to_vec`](@ref) is def
 function j′vp(fdm, f, ȳ, x)
     x_vec, vec_to_x = to_vec(x)
     ȳ_vec, _ = to_vec(ȳ)
-    return vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec))
+    return (vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec)), )
 end
-j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)
+j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)[1]
 
 """
     to_vec(x) -> Tuple{<:AbstractVector, <:Function}

test/grad.jl

@@ -20,14 +20,14 @@ Base.length(x::DummyType) = size(x.X, 1)
         rng, fdm = MersenneTwister(123456), central_fdm(5, 1)
         x = randn(rng, T, 2)
         xc = copy(x)
-        @test grad(fdm, x->sin(x[1]) + cos(x[2]), x) ≈ [cos(x[1]), -sin(x[2])]
+        @test grad(fdm, x->sin(x[1]) + cos(x[2]), x)[1] ≈ [cos(x[1]), -sin(x[2])]
         @test xc == x
     end
 
     function check_jac_and_jvp_and_j′vp(fdm, f, ȳ, x, ẋ, J_exact)
         xc = copy(x)
-        @test jacobian(fdm, f, x; len=length(ȳ)) ≈ J_exact
-        @test jacobian(fdm, f, x) == jacobian(fdm, f, x; len=length(ȳ))
+        @test jacobian(fdm, f, x; len=length(ȳ))[1] ≈ J_exact
+        @test jacobian(fdm, f, x)[1] == jacobian(fdm, f, x; len=length(ȳ))[1]
         @test _jvp(fdm, f, x, ẋ) ≈ J_exact * ẋ
         @test _j′vp(fdm, f, ȳ, x) ≈ transpose(J_exact) * ȳ
         @test xc == x
@@ -56,46 +56,46 @@ Base.length(x::DummyType) = size(x.X, 1)
             @testset "check multiple matrices" begin
                 x, y = rand(rng, 3, 3), rand(rng, 3, 3)
                 jac_xs = jacobian(fdm, f1, x, y)
-                @test jac_xs[1] ≈ jacobian(fdm, x->f1(x, y), x)
-                @test jac_xs[2] ≈ jacobian(fdm, y->f1(x, y), y)
+                @test jac_xs[1] ≈ jacobian(fdm, x->f1(x, y), x)[1]
+                @test jac_xs[2] ≈ jacobian(fdm, y->f1(x, y), y)[1]
             end
 
             @testset "check mixed scalar and matrices" begin
                 x, y = rand(3, 3), 2
                 jac_xs = jacobian(fdm, f1, x, y)
-                @test jac_xs[1] ≈ jacobian(fdm, x->f1(x, y), x)
-                @test jac_xs[2] ≈ jacobian(fdm, y->f1(x, y), y)
+                @test jac_xs[1] ≈ jacobian(fdm, x->f1(x, y), x)[1]
+                @test jac_xs[2] ≈ jacobian(fdm, y->f1(x, y), y)[1]
             end
         end
 
         @testset "grad" begin
             @testset "check multiple matrices" begin
                 x, y = rand(rng, 3, 3), rand(rng, 3, 3)
                 dxs = grad(fdm, f2, x, y)
-                @test dxs[1] ≈ grad(fdm, x->f2(x, y), x)
-                @test dxs[2] ≈ grad(fdm, y->f2(x, y), y)
+                @test dxs[1] ≈ grad(fdm, x->f2(x, y), x)[1]
+                @test dxs[2] ≈ grad(fdm, y->f2(x, y), y)[1]
             end
 
             @testset "check mixed scalar & matrices" begin
                 x, y = rand(rng, 3, 3), 2
                 dxs = grad(fdm, f2, x, y)
-                @test dxs[1] ≈ grad(fdm, x->f2(x, y), x)
-                @test dxs[2] ≈ grad(fdm, y->f2(x, y), y)        
+                @test dxs[1] ≈ grad(fdm, x->f2(x, y), x)[1]
+                @test dxs[2] ≈ grad(fdm, y->f2(x, y), y)[1]
             end
 
             @testset "check tuple" begin
                 x, y = rand(rng, 3, 3), 2
-                dxs = grad(fdm, f3, (x, y))
-                @test dxs[1] ≈ grad(fdm, x->f3((x, y)), x)
-                @test dxs[2] ≈ grad(fdm, y->f3((x, y)), y)    
+                dxs = grad(fdm, f3, (x, y))[1]
+                @test dxs[1] ≈ grad(fdm, x->f3((x, y)), x)[1]
+                @test dxs[2] ≈ grad(fdm, y->f3((x, y)), y)[1]   
             end
 
             @testset "check dict" begin
                 x, y = rand(rng, 3, 3), 2
                 d = Dict(:x=>x, :y=>y)
-                dxs = grad(fdm, f4, d)
-                @test dxs[:x] ≈ grad(fdm, x->f3((x, y)), x)
-                @test dxs[:y] ≈ grad(fdm, y->f3((x, y)), y)
+                dxs = grad(fdm, f4, d)[1]
+                @test dxs[:x] ≈ grad(fdm, x->f3((x, y)), x)[1]
+                @test dxs[:y] ≈ grad(fdm, y->f3((x, y)), y)[1]
             end
         end
     end
@@ -168,8 +168,8 @@ Base.length(x::DummyType) = size(x.X, 1)
         x, y = randn(rng, T, N), randn(rng, T, M)
         z̄ = randn(rng, T, N + M)
         xy = vcat(x, y)
-        x̄ȳ_manual = j′vp(fdm, xy->sin.(xy), z̄, xy)
-        x̄ȳ_auto = j′vp(fdm, x->sin.(vcat(x[1], x[2])), z̄, (x, y))
+        x̄ȳ_manual = j′vp(fdm, xy->sin.(xy), z̄, xy)[1]
+        x̄ȳ_auto = j′vp(fdm, x->sin.(vcat(x[1], x[2])), z̄, (x, y))[1]
         x̄ȳ_multi = j′vp(fdm, (x, y)->sin.(vcat(x, y)), z̄, x, y)
         @test x̄ȳ_manual ≈ vcat(x̄ȳ_auto...)
         @test x̄ȳ_manual ≈ vcat(x̄ȳ_multi...)