Interpolations, between Lux Dense layers, supported?

General Usage
, , ,
1

Greetings, one of the custom Lux models I am trying to develop requires interpolation in between, something like the following(?) which errs

using DataInterpolations, Lux, Random, Zygote, Enzyme

function convert_interpolation_then_vector(M :: AbstractMatrix)
    vectorized_data = collect.(eachslice(M,dims=1))
    vec = LagrangeInterpolation(vectorized_data[1], time)
    return vec'
end

temp_model = Chain(Dense(1, 1, use_bias = false), convert_interpolation_then_vector, Dense(1,1))



ERROR: Need an adjoint for constructor QuadraticInterpolation{Vector{Float64}, Vector{Float64}, true, Float64}. Gradient is of type Vector{Float64}

And gradient calculations for some of the interpolations work on Enzyme but not on Zygote.

time = collect(0:0.01:1)
x = LinearInterpolation(sin.(time), time)
Enzyme.gradient(Forward, time -> sum(x(time)), time) #works

Zygote.gradient(time -> sum(x(time)), time) # Mutating Array error

But I am not sure if Enzyme will work on Lux models currently for gradient calculations? Can anyone suggest what I am missing? Thank you very much for your help.

2

You can always make a custom layer for Lux, then write an rrule that uses Enzyme which will hide the Enzyme use from Zygote.

If it’s not clear what I mean by that - look at ChainRules doc, but then instead of manually implementing the rrule you can call Enzyme.

3

@cortner’s solution is the best in the short term.

However, for most CPU Lux models, Enzyme should be fine if your code is type-stable and such. If it doesn’t work, open an issue.