EnzymeNoTypeError when diff'ing a loss function with FFTA

I am trying to minimize a Lux network with LBFGS. My loss function contains an fft transform, because I'm trying to predict spectra and I want to constrain in both domains.
I've been able to get this working fine with Reactant and Lux and just using an Adam optimizer, but when I try to switch to LBFGS and using Enzyme, I'm getting the error thats's shown below the MWE (and which suggests filing a bug report). I tried FFTW before, but that just gave me a different error
```
using NeuralOperators
using Lux
using Random
using Reactant
using FFTA
using StatsBase: mean

using ComponentArrays
using Enzyme
using Optimization
using OptimizationOptimJL
using Statistics

Enzyme.@import_rrule(typeof(*), AbstractFFTs.Plan, AbstractArray)
EnzymeRules.inactive(::typeof(plan_fft), args...) = nothing

rng = Random.default_rng()
Random.seed!(rng, 1234)

xdev = reactant_device()
cdev = cpu_device()
const N_PADDED = 2048

function fft_loss(ps, p)
model, input, truth, st = p
preds, _ = model(input, ps, st)
k_loss_mse = mean(abs2, preds .- truth)

n_samples = size(truth, 2)
padded_truth = zeros(eltype(truth), N_PADDED, n_samples)
padded_truth[1:size(truth, 1), :] .= truth
padded_preds = zeros(eltype(preds), N_PADDED, n_samples)
padded_preds[1:size(preds, 1), :] .= preds

yᵣ = rfft(padded_truth, 1)
ŷᵣ = rfft(padded_preds, 1)
r_real_loss_mse = mean(abs2, real.(ŷᵣ) .- real.(yᵣ))
r_imag_loss_mse = mean(abs2, imag.(ŷᵣ) .- imag.(yᵣ))

total_loss = k_loss_mse + r_real_loss_mse + r_imag_loss_mse
return total_loss
end

function train!(model, ps, st, data; loss_function)
x_train, y_train = data
ps_ca = ComponentArray(ps)
adtype = Optimization.AutoEnzyme(mode = Enzyme.set_runtime_activity(Enzyme.Reverse))
optf = OptimizationFunction(loss_function, adtype)
prob = OptimizationProblem(optf, ps_ca, (model, x_train, y_train, st))

opt = LBFGS(m = 10)
sol = solve(prob, opt, maxiters = 1000)
return sol.u, st, losses
end

function main()
kAxis = range(1.0f0, 10.0f0, step = 0.02f0)
EVAL_POINTS = length(kAxis)
N_SENSORS = 450
u_data = rand(Float32, N_SENSORS, 5)
v_data = rand(Float32, EVAL_POINTS, 5)
y_data = Float32.(collect(reshape(kAxis, 1, EVAL_POINTS)))
data = ((u_data, y_data), v_data)

deeponet = DeepONet(
Chain(
Dense(N_SENSORS => 64, tanh; init_weight = glorot_uniform(Float32))
),
Chain(
Dense(1 => 64, tanh; init_weight = glorot_uniform(Float32))
)
)
ps, st = Lux.setup(rng, deeponet)
return train!(deeponet, ps, st, data; loss_function = fft_loss)
end

main()
```

Error message:
```
┌ Warning: TODO forward zero-set of memorycopy used memset rather than runtime type
│ Caused by:
│ Stacktrace:
│ [1] copy
│ @ ./array.jl:350
│ [2] unaliascopy
│ @ ./abstractarray.jl:1516
│ [3] unaliascopy
│ @ /qfs/people/stru821/.julia/juliaup/julia-1.11.8+0.x64.linux.gnu/share/julia/stdlib/v1.11/LinearAlgebra/src/adjtrans.jl:92
│ [4] unalias
│ @ ./abstractarray.jl:1500
│ [5] broadcast_unalias
│ @ ./broadcast.jl:946
│ [6] preprocess
│ @ ./broadcast.jl:953
│ [7] preprocess_args
│ @ ./broadcast.jl:956
│ [8] preprocess
│ @ ./broadcast.jl:952
│ [9] copyto!
│ @ ./broadcast.jl:969
│ [10] copyto!
│ @ ./broadcast.jl:925
│ [11] materialize!
│ @ ./broadcast.jl:883
│ [12] materialize!
│ @ ./broadcast.jl:880
│ [13] materialize!
│ @ ./broadcast.jl:0
└ @ Enzyme.Compiler ~/.julia/packages/Enzyme/OUOuC/src/rules/llvmrules.jl:814
ERROR: LoadError: EnzymeNoTypeError: Enzyme cannot statically prove the type of a value being differentiated and risks a correctness error if it gets it wrong.
Generally this shouldn't occur as Enzyme records type information from julia, but may be expected if you, for example, copy untyped data.
or alternatively emit very large sized registers that exceed the maximum size of Enzyme's type analysis. If it seems reasonable to differentiate
this code, open an issue! If the cause of the error is too large of a register, you can request Enzyme increase the size (https://enzyme.mit.edu/julia/dev/api/#Enzyme.API.maxtypeoffset!-Tuple{Any})
or depth (https://enzyme.mit.edu/julia/dev/api/#Enzyme.API.maxtypedepth!-Tuple{Any}) of its type analysis.
Alternatively, you can tell Enzyme to take its best guess from context with (https://enzyme.mit.edu/julia/dev/api/#Enzyme.API.looseTypeAnalysis!-Tuple{Any})
All of these settings are global configurations that need to be set immediately after loading Enzyme, before any differentiation occurs.
To toggle more information for debugging (needed for bug reports), set Enzyme.Compiler.VERBOSE_ERRORS[] = true (default false)
Failure within method: MethodInstance for FFTA.CallGraphNode!(::Vector{FFTA.CallGraphNode{ComplexF32}}, ::Int64, ::Vector{Vector{ComplexF32}}, ::Int64, ::Int64)
Hint: catch this exception as `err` and call `code_typed(err)` to inspect the errornous code.

Stacktrace:
[1] __safe_setindex!
@ ./array.jl:1001 [inlined]
[2] push!
@ ./array.jl:1262 [inlined]
[3] CallGraphNode!
@ ~/.julia/packages/FFTA/fhqKV/src/callgraph.jl:77
[4] CallGraph
@ ~/.julia/packages/FFTA/fhqKV/src/callgraph.jl:124 [inlined]
[5] #plan_rfft#7
@ ~/.julia/packages/FFTA/fhqKV/src/plan.jl:87 [inlined]
[6] plan_rfft
@ ~/.julia/packages/FFTA/fhqKV/src/plan.jl:79 [inlined]
[7] rfft
@ ~/.julia/packages/AbstractFFTs/4iQz5/src/definitions.jl:67
[8] fft_loss
@ /qfs/projects/bioprep/users/stru821/JuliaML/Lux_Examples/Bfgs_test.jl:35
[9] firstapply
@ ~/.julia/packages/OptimizationBase/mYxHK/ext/OptimizationEnzymeExt.jl:12 [inlined]
[10] firstapply
@ ~/.julia/packages/OptimizationBase/mYxHK/ext/OptimizationEnzymeExt.jl:0 [inlined]
[11] augmented_julia_firstapply_9820_inner_36wrap
@ ~/.julia/packages/OptimizationBase/mYxHK/ext/OptimizationEnzymeExt.jl:0
[12] macro expansion
@ ~/.julia/packages/Enzyme/OUOuC/src/compiler.jl:6672 [inlined]
[13] enzyme_call
@ ~/.julia/packages/Enzyme/OUOuC/src/compiler.jl:6151 [inlined]
[14] AugmentedForwardThunk
@ ~/.julia/packages/Enzyme/OUOuC/src/compiler.jl:6099 [inlined]
[15] autodiff
@ ~/.julia/packages/Enzyme/OUOuC/src/Enzyme.jl:419 [inlined]
[16] (::OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}})(res::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, θ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, p::Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}})
@ OptimizationEnzymeExt ~/.julia/packages/OptimizationBase/mYxHK/ext/OptimizationEnzymeExt.jl:121
[17] grad
@ ~/.julia/packages/OptimizationBase/mYxHK/ext/OptimizationEnzymeExt.jl:120 [inlined]
[18] (::OptimizationOptimJL.var"#10#16"{OptimizationCache{LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Bool, OptimizationOptimJL.var"#6#8", Nothing}, OptimizationOptimJL.var"#9#15"{OptimizationCache{LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Bool, OptimizationOptimJL.var"#6#8", Nothing}}})(G::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, θ::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}})
@ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/VKmZA/src/OptimizationOptimJL.jl:198
[19] value_gradient!!(obj::TwiceDifferentiable{Float32, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Float32, ComponentMatrix{Float32, Matrix{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}, Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}}, x::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}})
@ NLSolversBase ~/.julia/packages/NLSolversBase/gOqwS/src/interface.jl:92
[20] value_gradient!(obj::TwiceDifferentiable{Float32, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Float32, ComponentMatrix{Float32, Matrix{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}, Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}}, x::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}})
@ NLSolversBase ~/.julia/packages/NLSolversBase/gOqwS/src/interface.jl:82
[21] initial_state(method::LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#14"{OptimizationCache{LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Bool, OptimizationOptimJL.var"#6#8", Nothing}, Base.RefValue{Int64}}}, d::TwiceDifferentiable{Float32, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Float32, ComponentMatrix{Float32, Matrix{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}, Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}}, x0::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}})
@ Optim ~/.julia/packages/Optim/mv9zc/src/multivariate/solvers/first_order/l_bfgs.jl:177
[22] optimize(d::TwiceDifferentiable{Float32, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Float32, ComponentMatrix{Float32, Matrix{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}, Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}}, initial_x::ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, method::LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, options::Optim.Options{Float64, OptimizationOptimJL.var"#_cb#14"{OptimizationCache{LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Bool, OptimizationOptimJL.var"#6#8", Nothing}, Base.RefValue{Int64}}})
@ Optim ~/.julia/packages/Optim/mv9zc/src/multivariate/optimize/optimize.jl:45
[23] __solve(cache::OptimizationCache{LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Bool, OptimizationOptimJL.var"#6#8", Nothing})
@ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/VKmZA/src/OptimizationOptimJL.jl:258
[24] solve!(cache::OptimizationCache{LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}, true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), OptimizationEnzymeExt.var"#grad#26"{Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ReverseMode{false, true, false, FFIABI, false, false}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, OptimizationBase.ReInitCache{ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Bool, OptimizationOptimJL.var"#6#8", Nothing})
@ OptimizationBase ~/.julia/packages/OptimizationBase/mYxHK/src/solve.jl:216
[25] solve(::OptimizationProblem{true, OptimizationFunction{true, AutoEnzyme{ReverseMode{false, true, false, FFIABI, false, false}, Nothing}, typeof(fft_loss), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, ComponentVector{Float32, Vector{Float32}, Tuple{Axis{(layer_1 = ViewAxis(1:28992, Axis(branch = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(layer_1 = ViewAxis(1:28864, Axis(weight = ViewAxis(1:28800, ShapedAxis((64, 450))), bias = ViewAxis(28801:28864, Shaped1DAxis((64,))))),)), layer_2 = ViewAxis(28865:28864, Shaped1DAxis((0,))))), trunk = ViewAxis(28865:28992, Axis(layer_1 = ViewAxis(1:128, Axis(weight = ViewAxis(1:64, ShapedAxis((64, 1))), bias = ViewAxis(65:128, Shaped1DAxis((64,))))),)))), layer_2 = ViewAxis(28993:28992, Shaped1DAxis((0,))))}}}, Tuple{DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}, @NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, @Kwargs{}}, ::LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Returns{Nothing}}; kwargs::@Kwargs{maxiters::Int64})
@ OptimizationBase ~/.julia/packages/OptimizationBase/mYxHK/src/solve.jl:98
[26] train!(model::DeepONet{Chain{@NamedTuple{layer_1::Parallel{typeof(*), @NamedTuple{branch::Chain{@NamedTuple{layer_1::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}, trunk::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, WeightInitializers.PartialFunction.Partial{Float32, typeof(glorot_uniform), Nothing, @Kwargs{}}, Nothing, Static.True}}, Nothing}}, Nothing}, layer_2::WrappedFunction{typeof(adjoint)}}, Nothing}}, ps::@NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}}, layer_2::@NamedTuple{}}, st::@NamedTuple{layer_1::@NamedTuple{branch::@NamedTuple{layer_1::@NamedTuple{layer_1::@NamedTuple{}}, layer_2::@NamedTuple{}}, trunk::@NamedTuple{layer_1::@NamedTuple{}}}, layer_2::@NamedTuple{}}, data::Tuple{Tuple{Matrix{Float32}, Matrix{Float32}}, Matrix{Float32}}; loss_function::Function)
@ Main /qfs/projects/bioprep/users/stru821/JuliaML/Lux_Examples/Bfgs_test.jl:52
[27] main()
@ Main /qfs/projects/bioprep/users/stru821/JuliaML/Lux_Examples/Bfgs_test.jl:74
[28] top-level scope
@ /qfs/projects/bioprep/users/stru821/JuliaML/Lux_Examples/Bfgs_test.jl:77
in expression starting at /qfs/projects/bioprep/users/stru821/JuliaML/Lux_Examples/Bfgs_test.jl:77
```