Basic question on square root example

[alexboche]

Hi there,

Nice package! I'm wondering if someone could suggest how to use this package to compute the derivative of a function that computes the square roots of a number by using a solver. I asked this question on discourse here.

Thanks

[gdalle]

Hey there @alexboche!
For technical reasons, the inputs and outputs have to be vectors, currently it fails for plain numbers.
Here is a working code sample:

julia> import ChainRulesCore, ImplicitDifferentiation, Zygote

julia> custom_sqrt(x) = sqrt.(x)  # plug in your own implementation
custom_sqrt (generic function with 1 method)

julia> ChainRulesCore.@non_differentiable custom_sqrt(x)  # make sure we don't autodiff through it

julia> conditions(x, y) = y .^ 2 .- x  # y = √x <=> y^2 = x
conditions (generic function with 1 method)

julia> differentiable_custom_sqrt = ImplicitDifferentiation.ImplicitFunction(custom_sqrt, conditions)
ImplicitDifferentiation.ImplicitFunction{typeof(custom_sqrt), typeof(conditions), typeof(Krylov.gmres)}(custom_sqrt, conditions, Krylov.gmres)

julia> x = rand(3)
3-element Vector{Float64}:
 0.8729398964128972
 0.17440935918324196
 0.030649123026444847

julia> Zygote.jacobian(x -> sqrt.(x), x)[1]
3×3 Matrix{Float64}:
 0.535153  0.0      0.0
 0.0       1.19725  0.0
 0.0       0.0      2.85602

julia> Zygote.jacobian(x -> custom_sqrt(x), x)[1]  # this function is not autodiff-compatible
3×3 Matrix{Float64}:
 0.0  0.0  0.0
 0.0  0.0  0.0
 0.0  0.0  0.0

julia> Zygote.jacobian(differentiable_custom_sqrt, x)[1]  # this one is, and the Jacobian is correct
3×3 Matrix{Float64}:
  0.535153  -0.0      -0.0
 -0.0        1.19725  -0.0
 -0.0       -0.0       2.85602

[alexboche]

I'd like to follow up on this. I'm now trying this on my actual use-case and having a bit of trouble. I try the following:

α = 1; c = [1.0, 2.0]; ξ = [1.0, 2.0]; ψ = 1.0

u(x) = -(1/ψ)*exp(-ψ*x)

function premsol(o)
    function f(prem)
        return prem[1] .- softmax(ξ + α*u.(-prem[1] .- o))'*(c-o)
    end
    prem0 = [0.0]
    res = nlsolve(f, prem0)
    prem = res.zero
    return prem[1]
end

conditions(o, prem) = prem .- softmax(ξ + α*u.(-prem .- o))'*(c-o)
differentiable_premsol = ImplicitDifferentiation.ImplicitFunction(premsol, conditions)

The function premsol seems to work fine, but I get an error when I try to differentiate differentiable_premsol with Zygote. Any thoughts on what's wrong? Thanks

Zygote.jacobian(x -> differentiable_premsol(x), x)
> 
MethodError: no method matching vec(::Float64)
Closest candidates are:
  vec(::LinearAlgebra.Transpose{<:Number, <:AbstractVector}) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/adjtrans.jl:215
  vec(::LinearAlgebra.Adjoint{<:Real, <:AbstractVector}) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/adjtrans.jl:216
  vec(::SparseArrays.AbstractSparseVector) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/SparseArrays/src/sparsevector.jl:964
  ...

Stacktrace:
  [1] (::ImplicitDifferentiation.var"#implicit_pullback#11"{Float64, Vector{Float64}, LinearOperators.LinearOperator{Float64, Int64, ImplicitDifferentiation.var"#mul_Bᵀ!#10"{ComposedFunction{typeof(last), Zygote.var"#ad_pullback#50"{Tuple{ImplicitDifferentiation.var"#conditions_x#7"{Float64, typeof(conditions)}, Vector{Float64}}, typeof(∂(λ))}}, Float64}, Nothing, Nothing, Vector{Float64}}, LinearOperators.LinearOperator{Float64, Int64, ImplicitDifferentiation.var"#mul_Aᵀ!#9"{ComposedFunction{typeof(last), Zygote.var"#ad_pullback#50"{Tuple{ImplicitDifferentiation.var"#conditions_y#8"{Vector{Float64}, typeof(conditions)}, Float64}, typeof(∂(λ))}}, Float64}, Nothing, Nothing, Vector{Float64}}, typeof(Krylov.gmres)})(dy::Float64)
    @ ImplicitDifferentiation ~/.julia/packages/ImplicitDifferentiation/90Hcu/src/implicit_function.jl:107
  [2] ZBack
    @ ~/.julia/packages/Zygote/D7j8v/src/compiler/chainrules.jl:205 [inlined]
  [3] Pullback
    @ ./In[24]:1 [inlined]
  [4] (::typeof(∂(#26)))(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface2.jl:0
  [5] (::Zygote.var"#206#207"{Tuple{Tuple{Nothing}}, typeof(∂(#26))})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/lib/lib.jl:205
  [6] (::Zygote.var"#1894#back#208"{Zygote.var"#206#207"{Tuple{Tuple{Nothing}}, typeof(∂(#26))}})(Δ::Float64)
    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67
  [7] Pullback
    @ ./operators.jl:1085 [inlined]
  [8] (::typeof(∂(#_#83)))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface2.jl:0
  [9] (::Zygote.var"#206#207"{Tuple{Tuple{Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(#_#83))})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/lib/lib.jl:205
 [10] #1894#back
    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]
 [11] Pullback
    @ ./operators.jl:1085 [inlined]
 [12] (::typeof(∂(ComposedFunction{typeof(Zygote._jvec), var"#26#27"}(Zygote._jvec, var"#26#27"()))))(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface2.jl:0
 [13] (::Zygote.var"#60#61"{typeof(∂(ComposedFunction{typeof(Zygote._jvec), var"#26#27"}(Zygote._jvec, var"#26#27"())))})(Δ::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/compiler/interface.jl:41
 [14] withjacobian(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/lib/grad.jl:150
 [15] jacobian(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/D7j8v/src/lib/grad.jl:128
 [16] top-level scope
    @ In[24]:1
 [17] eval
    @ ./boot.jl:373 [inlined]
 [18] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String)
    @ Base ./loading.jl:1196

[gdalle]

Hi @alexboche! It seems to me that your premsol function returns a scalar. But ImplicitDifferentation only works with vector inputs and outputs. If you make the appropriate change, does that solve your issue?

[alexboche]

That worked thanks!

[alexboche]

Hi @gdalle. I'm struggling with another use of this. Not sure what I'm doing wrong--would you mind looking at my code?
I'm guessing I'm using scalars somewhere that I should be using vectors as above, but I'm gonna have to access the scalar components at some point, and I believe the key function recov_pi_s1, which I'm trying to autodiff through, does take and return a vector.

Any thoughts? Thanks!

Background code:

using NLsolve
using ImplicitDifferentiation
using Zygote

ψ = 3
u(x) = -(1/ψ)*exp(-ψ*x)
c1 = .8
c2 = .3

function act(pi,s1,o1,o21,o22)
   (1-s1)*(o22)*(c1-o1) + s1*(1-o21)*(c2-o21) + (1-s1)*(1-o22)*(c1 + c2 - o1 - o22)  -  pi
end


function thresh(pi,s1,o1,o21,o22) # equal zero # ex-ante ic
    (1/ψ)*(u(-pi - o1 - o22) - u(-pi - o1))  +  (1-o22)*u(-pi - o1 - o22)   -  # eu if get at threshold 
    ( (1/ψ)*(u(-pi - s1 - o21) - u(-pi - s1 - 0))  +      # eu if don't get at thresh
    (1-o21)*u(-pi - s1 - o21) )
end

act_thresh(pi,s1,o1,o20,o21) = [act(pi,s1,o1,o20,o21), thresh(pi,s1,o1,o20,o21)] # helper

Main code:

function recov_pi_s1(arg_vec) # key function 
    # Returns [pi,s1]
    o1,o20,o21 = arg_vec
    x0 = [0.0,0.0]
    ans = nlsolve(x -> act_thresh(x[1],x[2],o1,o20,o21), x0).zero
    return ans[1:2]
end

conditions(x,y) = act_thresh(y...,x...)  # y is [pi,s1], x is [o1,o20,o21]

diff_recov_pi_s1 = ImplicitDifferentiation.ImplicitFunction(recov_pi_s1, conditions)

x0 = [.5,.8,.5]

recov_pi_s1(x0) # works
diff_recov_pi_s1(x0) # works # ->
 # 2-element Vector{Float64}:
#  0.0988832449553662
#  0.33705585014877953

jacobian(w -> diff_recov_pi_s1(w), x0) # doesn't work
> 
MethodError: no method matching vec(::Tuple{Float64, Float64})
Closest candidates are:
  vec(::LinearAlgebra.Transpose{<:Number, <:AbstractVector}) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/adjtrans.jl:215
  vec(::LinearAlgebra.Adjoint{<:Real, <:AbstractVector}) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/adjtrans.jl:216
  vec(::SparseArrays.AbstractSparseVector) at /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/SparseArrays/src/sparsevector.jl:964

Stacktrace: here

[gdalle]

OK so I'm not sure where it comes from but there is a problem with tuple destructuring.
I've played around a bit, and if you write your conditions like this instead, it should work:

function conditions(x,y)
    pi, s1 = y
    o1, o20, o21 = x
    return act_thresh(pi,s1,o1,o20,o21)
end

Can you confirm? Maybe open an issue with a simple MWE if you feel like it?

[alexboche]

That worked! I will have go back through this later today and try to understand the issue better. Sure, if there are further problems, I'll make a MWE in a separate issue. Thanks!