Compatability with Enzyme?

[wsmoses]

Hi,

Very fun package!

Any plans to add compatibility with enzyme?

[FriesischScott]

We have plans to use DifferentiationInterface to implement derivatives for all models in #304. We'll most likely default to Mooncake for native Julia models and FiniteDifferences for models that call third-party-solvers.

However, due to the way we use DataFrames internally I don't think Enzyme will work here. Unless you think it will?

[wsmoses]

I don't think theres any reason Enzyme would fail for dataframes offhand?

[FriesischScott]

I don't think theres any reason Enzyme would fail for dataframes offhand?

I'll make some experiments.

[FriesischScott]

@wsmoses This is is a MWE of how we'll likely implement the derivatives.

using UncertaintyQuantification
using ADTypes
using Mooncake
using DifferentiationInterface
using DataFrames
using Enzyme

m = Model(df -> df.x .^ 2 .* df.y, :z)

function f(x)
    df = DataFrame(; x=x[1], y=x[2])
    evaluate!(m, df)
    return only(df.z)
end

x = [0.5, 2.0]

value_and_gradient(f, AutoMooncake(), x)# (0.5, [2.0, 0.25])

Basically we'd have a wrapper function that we'll pass do DIfferentiationInterface which internally creates the DataFrame required for evaluation of the Model and returns the result. Using Mooncake this works perfectly.

Enzyme fails with an IllegalTypeAnalysisException

value_and_gradient(f, AutoEnzyme(), x)

and suggests to set Enzyme.API.strictAliasing!(false). When I do it doesn't error but also never finishes computing the gradient.

[gdalle]

If you want to help Enzyme devs help you, I suggest you try to reduce the example to remove as many external dependencies as possible. Getting rid of DI is relatively easy here since Enzyme has a gradient API you can use, but the more you simplify, the faster you can get a fix! I found LLMs to be rather good at creating Enzyme MWEs

[FriesischScott]

If you want to help Enzyme devs help you, I suggest you try to reduce the example to remove as many external dependencies as possible. Getting rid of DI is relatively easy here since Enzyme has a gradient API you can use, but the more you simplify, the faster you can get a fix! I found LLMs to be rather good at creating Enzyme MWEs

Thanks, I adapted the MWE to just use DataFrames and it results in the same error as expected.

using DataFrames
using Enzyme

function m(df)
    return df.x .^ 2 .* df.y
end

function f(x)
    df = DataFrame(; x=x[1], y=x[2])
    return m(df)
end

x = [0.5, 2.0]

Enzyme.gradient(Reverse, f, x)