Type instabilities?

[gdalle]

See this Discourse thread: https://discourse.julialang.org/t/custom-rule-for-differentiating-through-newton-solver-using-forwarddiff-works-for-gradient-fails-for-hessian/93002/20

[thorek1]

I think I encountered something related. the ForwardDiff extension flattens arrays to a vector on output (line 42 in ImplicitDifferentiationForwardDiffExt.jl) having the dimensions preserved should fix it in my case at least

[gdalle]

Unfortunately this seems required for the linear solvers to work: Krylov.jl doesn't work with arbitrary arrays. I tried to use KrylovKit.jl instead a while back and it ended up being type-unstable for different reasons, but I'm willing to try again

[gdalle]

Note that the type stability tests which used to be performed by ChainRules were removed temporarily because I bumped into JuliaDiff/ChainRulesTestUtils.jl#232

[thorek1]

I see, but can't you just reshape after the solve?

[gdalle]

I see, but can't you just reshape after the solve?

Not sure what you mean, I think this line does the flattening
https://github.com/gdalle/ImplicitDifferentiation.jl/blob/c6c8fc251b2b35d61497b35fb8f24e0b49d0ba0c/ext/ImplicitDifferentiationForwardDiffExt.jl#L35
and this line does the reshaping
https://github.com/gdalle/ImplicitDifferentiation.jl/blob/c6c8fc251b2b35d61497b35fb8f24e0b49d0ba0c/ext/ImplicitDifferentiationForwardDiffExt.jl#L38

[gdalle]

@thorek1 I think your issue might have been solved by @YingboMa in #41

[gdalle]

And the issue from the Discourse thread will be taken care of by #47

[thorek1]

Thanks, I will have a look and let you know.

[thorek1]

it didn't resolve my issue. the MWE is a bit involved (the error occurs at line 94). maybe a high level hint: Zygote works, ForwardDiff doesn't. ForwardDiff breaks when I try to use the result returned from implicit diff in computations because the dimension info is lost and he returns a vector instead of a matrix. the function I try to diff has 1 arg and 2 kwargs and returns a matrix and a bool. the conditions accepts three args, the same kwargs and returns a matrix of the same dimension as the first output form the forward function. in principle it should work and it does for Zygote

[gdalle]

OK I think I know where this is coming from.

[thorek1]

That fixed it. Thanks!