Error with LBFGS() with using adjoint on 3D arrays

[roflmaostc]

3D optimization with LBFGS fails. 2D arrays work because adjoint is probably defined.

# ╔═╡ 350d641c-ae24-11ee-018a-b331fe41f5b3
using Optimization, OptimizationOptimisers, OptimizationOptimJL, Zygote

# ╔═╡ 4c272f03-4aae-4788-99bd-88e8bb3a8eea
f(x, p) = sum(abs2, x)

# ╔═╡ eba0efac-2874-4031-97fc-3256ad103bd6
init0 = ones((10, 10, 10))

# ╔═╡ f0afb848-f40f-4999-86cc-50f63f7b553d
opt_fun = OptimizationFunction(f, AutoZygote())

# ╔═╡ c3171aae-9822-4add-a2ec-d1e749158684
problem = OptimizationProblem(opt_fun, init0);

# ╔═╡ d5d19f34-3100-43a3-8d35-6aae4bf1b171
res = solve(problem, OptimizationOptimJL.LBFGS(), maxiters=20)

results in:

adjoint not defined for Array{Float64, 3}. Consider using `permutedims` for higher-dimensional arrays.

    error(::String)@error.jl:35
    adjoint(::Array{Float64, 3})@transpose.jl:3
    alloc_H(::Array{Float64, 3}, ::Float64)@abstract.jl:25
    __solve(::Optimization.OptimizationCache{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), OptimizationZygoteExt.var"#38#56"{OptimizationZygoteExt.var"#37#55"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}}, OptimizationZygoteExt.var"#41#59"{OptimizationZygoteExt.var"#37#55"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}}, OptimizationZygoteExt.var"#45#63", Nothing, OptimizationZygoteExt.var"#49#67"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}, OptimizationZygoteExt.var"#53#71"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}, Nothing, Nothing, Nothing, Nothing, Nothing, Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Optim.var"#19#21"}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}, Bool, OptimizationOptimJL.var"#4#6"})@OptimizationOptimJL.jl:191
    solve!(::Optimization.OptimizationCache{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), OptimizationZygoteExt.var"#38#56"{OptimizationZygoteExt.var"#37#55"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}}, OptimizationZygoteExt.var"#41#59"{OptimizationZygoteExt.var"#37#55"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}}, OptimizationZygoteExt.var"#45#63", Nothing, OptimizationZygoteExt.var"#49#67"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}, OptimizationZygoteExt.var"#53#71"{SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}}, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Optimization.ReInitCache{Array{Float64, 3}, SciMLBase.NullParameters}, Nothing, Nothing, Nothing, Nothing, Nothing, Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Optim.var"#19#21"}, Base.Iterators.Cycle{Tuple{Optimization.NullData}}, Bool, OptimizationOptimJL.var"#4#6"})@solve.jl:177
    var"#solve#599"(::Base.Pairs{Symbol, Int64, Tuple{Symbol}, NamedTuple{(:maxiters,), Tuple{Int64}}}, ::typeof(CommonSolve.solve), ::SciMLBase.OptimizationProblem{true, SciMLBase.OptimizationFunction{true, ADTypes.AutoZygote, typeof(Main.var"workspace#9".f), Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED_NO_TIME), Nothing, Nothing, SymbolicIndexingInterface.SymbolCache{Nothing, Nothing, Nothing}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing}, Array{Float64, 3}, SciMLBase.NullParameters, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}}, ::Optim.LBFGS{Nothing, LineSearches.InitialStatic{Float64}, LineSearches.HagerZhang{Float64, Base.RefValue{Bool}}, Optim.var"#19#21"})@solve.jl:94
    top-level scope@[Local: 1](http://localhost:1234/edit?id=350d6414-ae24-11ee-39ce-75009143ebb4#)[inlined]

status `/tmp/jl_is47EO/Project.toml`
⌅ [7f7a1694] Optimization v3.20.2
⌅ [36348300] OptimizationOptimJL v0.1.14
⌅ [42dfb2eb] OptimizationOptimisers v0.1.6
  [e88e6eb3] Zygote v0.6.68
  [44cfe95a] Pkg v1.9.2
Info Packages marked with ⌅ have new versions available but compatibility constraints restrict them from upgrading. To see why use `status --outdated`


Julia Version 1.9.4
Commit 8e5136fa297 (2023-11-14 08:46 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 24 × AMD Ryzen 9 5900X 12-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-14.0.6 (ORCJIT, znver3)
  Threads: 24 on 24 virtual cores
Environment:
  JULIA_NUM_THREADS = 24
  JULIA_DEPOT_PATH = /home/felix/.julia:/home/felix/.julia/juliaup/julia-1.9.4+0.x64.linux.gnu/local/share/julia:/home/felix/.julia/juliaup/julia-1.9.4+0.x64.linux.gnu/share/julia
  JULIA_PROJECT = /home/felix/.julia/dev/SwissVAMyKnife.jl/Project.toml
  JULIA_LOAD_PATH = @:@v#.#:@stdlib
  JULIA_REVISE_WORKER_ONLY = 1

[Vaibhavdixit02]

This error is from Optim's alloc_H method which relies on being able to transpose the optimization variable. Replacing that with Adam actually works so it's not really an issue but more of a limitation of the methods

res = solve(problem, OptimizationOptimisers.Adam(), maxiters=20)

[roflmaostc]

I'm not sure if I understand but Optim.jl can optimize on 3D (ND even) arrays.

I raised this issue since I always use Optim.jl and it works but it failed with Optimization.jl

[Vaibhavdixit02]

Interesting, do you use out-of-place derivative evaluations there?

[roflmaostc]

Usually Optim.optimize(Optim.only_fg!(fg!), x, opt=LBFGS()) but all of it works.