Testing Differentiation Interfase and check efficient boost using Mooncake

[facusapienza21]

No description provided.

[facusapienza21]

I open this PR to check whether MoonCake can be use to compute the gradients more efficiently in SphereUDE. With this core, the benchmark for the following adjoint methods give the following results:

sensealg_types = [
    GaussAdjoint(autojacvec = ReverseDiffVJP(true)),
    InterpolatingAdjoint(autojacvec = ReverseDiffVJP(true)),
    QuadratureAdjoint(autojacvec = ReverseDiffVJP(true)),
    GaussAdjoint(autojacvec = MooncakeVJP()),
    InterpolatingAdjoint(autojacvec = MooncakeVJP()),
    QuadratureAdjoint(autojacvec = MooncakeVJP()),
    SphereBackSolveAdjoint()
]

I don't observe any incredible increase in efficiency due to Mooncake, but also I am not sure I am using right. All this benchmarks are done without any regularization, so nested AD is not really required.

[wsmoses]

out of curiosity does using Enzyme help?

[facusapienza21]

Thank you for jumping in @wsmoses !

Right now when I just use autojacvec = EnzymeVJP() the sensitivity analysis just does not work. I know Enzyme requires some more customization on how variables are defined, and I haven't followed that rabbit hole here (as @albangossard and @JordiBolibar did in the case of ODINN.jl).

Here, Mooncake does not seem to provide huge improvement (and I am right to assume Mooncake should be as fast as Enzyme @albangossard ?).

I also want to re-direct attention to another feature in SphereUDE which is the regularizations that right now are handled by Zygote in order to do nested differentiations (as explain in this post in the Lux documentation):

        if typeof(reg.diff_mode) <: LuxNestedAD


            if reg.diff_mode.method == "ForwardDiff"
                # This is giving now the wrong results!!!
                Jac = batched_jacobian(
                    smodel,
                    AutoForwardDiff(),
                    reshape(nodes, 1, n_quadrature),
                )
            elseif reg.diff_mode.method == "Zygote"
                # This can also be done with Zygote in reverse mode
                Jac = batched_jacobian(
                    smodel,
                    AutoZygote(),
                    reshape(nodes, 1, n_quadrature),
                )
            else
                throw("Method for AD backend no implemented.")
            end

It would be great to improve the performance of SphereUDE, since this is a very small UDE (so technically not very complicated) but that be used for many paleomagnetists, so the faster the tool the best! Any feedback, comments or contributions are more than welcome!!!