Repository: JuliaDiff/FiniteDifferences.jl
Pull request #51 (2019-10-04 22:39:59)

change return type of grad/jacobian to match args
I didn't notice this yesterday until I start working on the Zygote PR today. So the basically current return value type of single variable grad/jacobian is not very consistent with multivariate grad/jacobian. It should return `(x, )`, this is causing extra handling when I implement `gradcheck` using the jacobians in Zygote.

nickrobinson251: I will defer to WIll on this.

If this is what is mostly widely useful (and not surprising) behaviour in practice, then i'm okay with it. But we should be cobfident it is before making the change.

Also this is a breaking change, so at the very least it'd be good to bump the version number accordingly in this PR.

(As a complete aside I'm really looking forward to being able to unwrap singe-element collections with `only(f(x))` instead of `f(x)[1]` in Julia 1.4+)
codecov-io: # [Codecov](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=h1) Report
> Merging [#51](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=desc) into [master](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/commit/22e248a25cb1370b7885274acd6dc41a5fe2c002?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `90.9%`.

[![Impacted file tree graph](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51/graphs/tree.svg?width=650&token=fIoJjgefQ8&height=150&src=pr)](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master      #51   +/-   ##
=======================================
  Coverage   98.63%   98.63%           
=======================================
  Files           4        4           
  Lines         147      147           
=======================================
  Hits          145      145           
  Misses          2        2
```


| [Impacted Files](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/grad.jl](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51/diff?src=pr&el=tree#diff-c3JjL2dyYWQuamw=) | `96.96% <90.9%> (ø)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=footer). Last update [22e248a...864199a](https://codecov.io/gh/JuliaDiff/FiniteDifferences.jl/pull/51?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

Roger-luo: @nickrobinson251 I can change it to `x[]` instead of `x[1]` if you prefer that. But practically I think I need this for Zygote or I have to implement an extra dispatch in Zygote that deal with single input. But notice this is also not consistent with `Zygote`'s `gradient`.
nickrobinson251: > practically I think I need this for Zygote or I have to implement an extra dispatch in Zygote that deal with single input. But notice this is also not consistent with Zygote's gradient.

Seem like compelling reasons to me! And the code looks good. But I'd like @willtebbutt before going ahead.

(Also, yeah, `[1]` is good here. And as you say we won't be able to use `only`, unless we add it to Comat.jl or whatever, it's just that cases like this remind me that I wish we had it / glad we will one day.)
willtebbutt: I'm pro- this change, but we need to enforce it for all functions, including `jvp` and `j′vp`. So if you could make those changes as well, that would be great.
willtebbutt: @Roger-luo any chance you're able to finish this off? It should only require minor modifications to `jvp` and `j′vp` to make this work as intended, and the changes would be analogous to the ones that you've already made.

Would be great to get this merged so that we can get [this PR](https://github.com/FluxML/Zygote.jl/pull/235) merged and resolve [this](https://github.com/FluxML/Zygote.jl/issues/283) and other related issues.
Roger-luo: Yeah, sorry I was having mid term exams... I'll have time on this in the weekend I think. @willtebbutt 
willtebbutt: Fantastic, thanks @Roger-luo . Hope the mid terms went okay!
Roger-luo: @willtebbutt just applied this to `jvp` and `j'vp` I think, `jvp` always returns a single number so just `(x, )` would be sufficient?
willtebbutt: If @oxinabox is happy, I'm happy to merge.
oxinabox: I have not reviewed but I trust @willtebbutt who knows this code base much better than I do
willtebbutt: Actually, we definitely want `j′vp` to return a `Tuple`, but I don't think that we want `jvp` to since there's not need to handle multiple outputs in the same way. We should probably revert to a single output, rather than a `Tuple`. Apologies for the confusion!
Roger-luo: @willtebbutt is this good?
willtebbutt: Yup, good by me. Will merge once tests have passed.
Roger-luo: BTW I'll need a tag to make CI in Zygote pass when this is merged later I think
willtebbutt: Understood. Could you bump to 0.9.0 in this PR?

Patch:
diff --git a/Project.toml b/Project.toml
index 80707ec..3c5207e 100644
--- a/Project.toml
+++ b/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"
diff --git a/src/grad.jl b/src/grad.jl
index e7668d3..c302155 100644
--- a/src/grad.jl
+++ b/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}, ẋ::AV{<:Number}) = fdm(ε -> f(x .+ ε .* x
 
 Convenience function to compute `transpose(jacobian(f, x)) * ȳ`.
 """
-_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))) * ȳ
+_j′vp(fdm, f, ȳ::AV{<:Number}, x::Vector{<:Number}) = transpose(jacobian(fdm, f, x; len=length(ȳ))[1]) * ȳ
 
 """
     jvp(fdm, f, x, ẋ)
@@ -98,7 +98,7 @@ function jvp(fdm, f, (x, ẋ)::Tuple{Any, Any})
 end
 function jvp(fdm, f, xẋs::Tuple{Any, Any}...)
     x, ẋ = collect(zip(xẋs...))
-    return jvp(fdm, xs->f(xs...), (x, ẋ))
+    return jvp(fdm, xs->f(xs...)[1], (x, 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, ȳ, x)
     x_vec, vec_to_x = to_vec(x)
     ȳ_vec, _ = to_vec(ȳ)
-    return vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec))
+    return (vec_to_x(_j′vp(fdm, x_vec->to_vec(f(vec_to_x(x_vec)))[1], ȳ_vec, x_vec)), )
 end
-j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)
+j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)[1]
 
 """
     to_vec(x) -> Tuple{<:AbstractVector, <:Function}


Test patch:
diff --git a/test/grad.jl b/test/grad.jl
index 4c0890b..31632bf 100644
--- a/test/grad.jl
+++ b/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, ȳ, x, ẋ, J_exact)
         xc = copy(x)
-        @test jacobian(fdm, f, x; len=length(ȳ)) ≈ J_exact
-        @test jacobian(fdm, f, x) == jacobian(fdm, f, x; len=length(ȳ))
+        @test jacobian(fdm, f, x; len=length(ȳ))[1] ≈ J_exact
+        @test jacobian(fdm, f, x)[1] == jacobian(fdm, f, x; len=length(ȳ))[1]
         @test _jvp(fdm, f, x, ẋ) ≈ J_exact * ẋ
         @test _j′vp(fdm, f, ȳ, x) ≈ transpose(J_exact) * ȳ
         @test xc == x
@@ -56,15 +56,15 @@ 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
 
@@ -72,30 +72,30 @@ Base.length(x::DummyType) = size(x.X, 1)
             @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̄ȳ_manual = j′vp(fdm, xy->sin.(xy), z̄, xy)
-        x̄ȳ_auto = j′vp(fdm, x->sin.(vcat(x[1], x[2])), z̄, (x, y))
+        x̄ȳ_manual = j′vp(fdm, xy->sin.(xy), z̄, xy)[1]
+        x̄ȳ_auto = j′vp(fdm, x->sin.(vcat(x[1], x[2])), z̄, (x, y))[1]
         x̄ȳ_multi = j′vp(fdm, (x, y)->sin.(vcat(x, y)), z̄, x, y)
         @test x̄ȳ_manual ≈ vcat(x̄ȳ_auto...)
         @test x̄ȳ_manual ≈ vcat(x̄ȳ_multi...)