Open benjaminfaber opened 3 weeks ago
Indeed ADTypes has no option for choosing betweeen autodiff
and autodiff_deferred
with Enzyme, partly because this choice should become unnecessary in the medium term (according to @wsmoses).
In the meantime, DifferentiationInterface allows the switch internally with the function DifferentiationInterface.nested(backend)
.
You can try setting conditions_x_backend
and conditions_y_backend
to DifferentiationInterface.nested(AutoEnzyme())
in your inner ImplicitFunction
to automatically use autodiff_deferred
.
I should also note that at the moment, Enzyme is only supported in forward mode by ImplicitDifferentiation (see #35). If your input is very high-dimensional (how big are we talking?), you may be better served by a reverse mode backend. Can you maybe give a mathematical description of the real function(s) you're trying to differentiate?
That bug is a GC error not the result of deferred.
i fixed two such bugs in Julia proper for the 1.10.5 release see if that helps in not open a mwe issue on enzyme proper
Deferred isn’t needed here, since the custom rule used is calling autodiff. Deferred is only required if doing higher order AD, Eg doing autodiff of autodiff
I think higher order is what's being done here: the conditions c2
involve autodiff
of implicit_f1
Anything in a custom rule is considered top level (aka does not have the deferred problem).
Of course if the code does call autodiff of autodiff elsewhere there may be a deferred issue.
Nevertheless the error log is a GC issue not a deferred error so any deferred use won’t fix it.
On Thu, Aug 15, 2024 at 1:27 AM Guillaume Dalle @.***> wrote:
I think higher order is what's being done here: the conditions c2 involve autodiff of implicit_f1
— Reply to this email directly, view it on GitHub https://github.com/JuliaDecisionFocusedLearning/ImplicitDifferentiation.jl/issues/153#issuecomment-2290662698, or unsubscribe https://github.com/notifications/unsubscribe-auth/AAJTUXBZAMOPXWBJW273EY3ZRQ333AVCNFSM6AAAAABMRE5ZGSVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDEOJQGY3DENRZHA . You are receiving this because you were mentioned.Message ID: <JuliaDecisionFocusedLearning/ImplicitDifferentiation. @.***>
Thank you both, for the clarification on this, and being willing to put up with questions from a relative newcomer to AD. To answer your question @gdalle, the intended application is a plasma physics/fusion energy optimization problem where we can have up to 300+ parameters in our state vector. The gradient calculation needed is ultimately going to be the gradient of a scalar optimization target, so reverse mode will be the way we want to go. The problem I'm currently trying to tackle has three nested optimization problems (A, B, C), where B depends on A and C depends on both B and A. I can use NonlinearSolve and the NonlinearLeastSquaresProblem, however I was not having any luck calculating the sensitivities whereas I have been able to (albeit slowly) with ForwardDiff and ImplicitDifferentiation, so I've been pursuing that route. I may be doing things incorrectly with NonlinearSolve, however we still have the ultimate goal of wanting to use a legacy Fortran application to do the forward solve in the implicit problem, so we'll need ImplicitDifferentiation for that.
I'll see what happens with I use a reverse mode differentiator and if that has any issues.
If you have any mathematical formulation of the nested optimization problems that you're free to share, id love to take a look!
Sure, I'd be happy to give a little mathematical background on the problem. We are trying to optimize the shape of the magnetic field used in magnetic confinement fusion concepts, which can be represented efficiently in cylindrical coordinates with a non-standard Fourier transformation:
R(r,\theta,\zeta) = \sum\limits_{m,n}^{M,N} R_{m,n}^c(r)\cos\left(n\theta + n\zeta\right) + R_{m,n}^s \sin\left(m\theta + n\zeta\right),
Z(r,\theta,\zeta) = \sum\limits_{m,n}^{M,N} Z_{m,n}^c(r)\cos\left(n\theta + n\zeta\right) + Z_{m,n}^s \sin\left(m\theta + n\zeta\right),
\phi = -\zeta
Derived quantities, like the Jacobian of the coordinate transformation, also have the same Fourier representation. I'm using Optimization.jl to find the Fourier modes of the derived quantities that satisfy their constraint conditions. For example, Jacobian equation satisfies the following condition
\frac{\partial}{\partial\theta}\left(\frac{1 + \frac{\partial R}{\partial \zeta}^2 + \frac{\partial Z}{\partial \zeta}^2 + \iota \left(\frac{\partial R}{\partial \theta} \frac{\partial R}{\partial \zeta} + \frac{\partial Z}{\partial \theta}\frac{\partial Z}{\partial \zeta}\right)}{\sqrt{g}}\right) + \frac{\partial}{\partial\zeta}\left(\frac{\frac{\partial R}{\partial \theta}\frac{\partial R}{\partial \zeta} + \frac{\partial Z}{\partial \theta}\frac{\partial Z}{\partial \zeta} + \iota \left(\frac{\partial R}{\partial \theta} ^2 + \frac{\partial Z}{\partial \theta} ^ 2\right)}{\sqrt{g}}\right) = 0.
Other quantities then depend on this Jacobian, so I've found it straight forward to solve for the Fourier modes of those quantities as another optimization problem where I can write down the optimality condition easily to use ImplicitDifferentiation to compute sensitivities, and is where I ran into this problem initially.
All that being said, I've found a workaround is to solve for all of the dependent quantities simultaneously, so that I only need to make one call to compute the derivative of an ImplicitFunction
and don't need to currently deal with this problem.
Thanks, glad you found a workaround! We'll revisit this when Julia 1.10.5 comes out
I've run into an issue where I want to compute the gradient of an implicit function that itself depends on another implicit function. I can do the operation successfully with
FowardDiff
, however I would like to use a backend that is faster, likeEnzyme
. When I run the MWE below, I run into a segmentation fault, I think related to the fact that one needs useautodiff_deferred
, but it looks like ADTypes doesn't yet have a backend option forautodiff_deferred
. Should I open an issue there or is this something that can be implemented inImplicitDifferentiation
/DifferentiationInterface
?Tagging @just-walk