Fix type-inference for j′vp (#201)

* Avoid composing many functions

* Add type-inferribility checks

* Increment version number

Author: sethaxen

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteDifferences"
 uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
-version = "0.12.21"
+version = "0.12.22"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/grad.jl

@@ -70,7 +70,7 @@ Compute an adjoint with any types of arguments `x` for which [`to_vec`](@ref) is
 function j′vp(fdm, f, ȳ, x)
     x_vec, vec_to_x = to_vec(x)
     ȳ_vec, _ = to_vec(ȳ)
-    return (vec_to_x(_j′vp(fdm, first ∘ to_vec ∘ f ∘ vec_to_x, ȳ_vec, x_vec)), )
+    return (vec_to_x(_j′vp(fdm, x -> first(to_vec(f(vec_to_x(x)))), ȳ_vec, x_vec)), )
 end
 
 j′vp(fdm, f, ȳ, xs...) = j′vp(fdm, xs->f(xs...), ȳ, xs)[1]

test/grad.jl

@@ -146,9 +146,10 @@ using FiniteDifferences: grad, jacobian, _jvp, jvp, j′vp, _j′vp, to_vec
             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)[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)
+            # Type inference: https://github.com/JuliaDiff/FiniteDifferences.jl/issues/199
+            x̄ȳ_manual = @inferred(j′vp(fdm, xy->sin.(xy), z̄, xy))[1]
+            x̄ȳ_auto = @inferred(j′vp(fdm, x->sin.(vcat(x[1], x[2])), z̄, (x, y)))[1]
+            x̄ȳ_multi = @inferred(j′vp(fdm, (x, y)->sin.(vcat(x, y)), z̄, x, y))
             @test x̄ȳ_manual ≈ vcat(x̄ȳ_auto...)
             @test x̄ȳ_manual ≈ vcat(x̄ȳ_multi...)
         end