Optimization over 3d Array throws

[jariji]

Describe the bug 🐞

I can optimize over Array{Float64, 3} with Optim.jl but it fails with Optimization.jl.

Expected behavior

No error.

Minimal Reproducible Example 👇

using Optimization, Optim, OptimizationOptimJL, Zygote, DifferentiationInterface
julia> let k = 5
           f(Ψ) = sum(abs2, Ψ)
           f(Ψ, _) = f(Ψ)
           @show f(rand(5,5,k), nothing)
           x0 = rand(k,k,k)
           grad!(dΨ, Ψ) = copyto!(dΨ, Zygote.gradient(f, Ψ)[1])
           res = optimize(f, grad!, x0, LBFGS()) #ok 
           @show res
           optfunc = OptimizationFunction(f, Optimization.AutoZygote())
           problem = OptimizationProblem(optfunc, x0)
           solve(problem, BFGS()) # error
       end
f(rand(5, 5, k), nothing) = 39.36502715916964
res =  * Status: success

 * Candidate solution
    Final objective value:     0.000000e+00

 * Found with
    Algorithm:     L-BFGS

 * Convergence measures
    |x - x'|               = 9.99e-01 ≰ 0.0e+00
    |x - x'|/|x'|          = Inf ≰ 0.0e+00
    |f(x) - f(x')|         = 3.85e+01 ≰ 0.0e+00
    |f(x) - f(x')|/|f(x')| = Inf ≰ 0.0e+00
    |g(x)|                 = 0.00e+00 ≤ 1.0e-08

 * Work counters
    Seconds run:   0  (vs limit Inf)
    Iterations:    1
    f(x) calls:    3
    ∇f(x) calls:   3

Error & Stacktrace ⚠️

ERROR: adjoint not defined for Array{Float64, 3}. Consider using `permutedims` for higher-dimensional arrays.
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] adjoint(a::Array{Float64, 3})
   @ LinearAlgebra /nix/store/x044gz5qmcy25gm3cjil3kg0mv5jnf5p-julia-bin-1.10.3/share/julia/stdlib/v1.10/LinearAlgebra/src/transpose.jl:3
 [3] alloc_H(x::Array{Float64, 3}, F::Float64)
   @ NLSolversBase ~/.julia/packages/NLSolversBase/kavn7/src/objective_types/abstract.jl:25
 [4] __solve(cache::OptimizationCache{…})
   @ OptimizationOptimJL ~/.julia/packages/OptimizationOptimJL/hDX5k/src/OptimizationOptimJL.jl:206
 [5] solve!(cache::OptimizationCache{…})
   @ SciMLBase ~/.julia/packages/SciMLBase/tEuIM/src/solve.jl:186
 [6] solve(::OptimizationProblem{…}, ::BFGS{…}; kwargs::@Kwargs{})
   @ SciMLBase ~/.julia/packages/SciMLBase/tEuIM/src/solve.jl:94
 [7] solve(::OptimizationProblem{…}, ::BFGS{…})
   @ SciMLBase ~/.julia/packages/SciMLBase/tEuIM/src/solve.jl:91
 [8] top-level scope

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

(jl_jwfzDz) pkg> st
Status `/tmp/jl_jwfzDz/Project.toml`
  [a0c0ee7d] DifferentiationInterface v0.6.24
  [429524aa] Optim v1.10.0
  [7f7a1694] Optimization v4.0.5
  [36348300] OptimizationOptimJL v0.4.1
  [e88e6eb3] Zygote v0.6.73

• Output of using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)

julia> using Pkg; Pkg.status(; mode = PKGMODE_MANIFEST)
Status `/tmp/jl_jwfzDz/Manifest.toml`
  [47edcb42] ADTypes v1.11.0
  [621f4979] AbstractFFTs v1.5.0
  [1520ce14] AbstractTrees v0.4.5
  [7d9f7c33] Accessors v0.1.38
  [79e6a3ab] Adapt v4.1.1
  [4fba245c] ArrayInterface v7.17.1
⌅ [a9b6321e] Atomix v0.1.0
  [fa961155] CEnum v0.5.0
  [082447d4] ChainRules v1.72.1
  [d360d2e6] ChainRulesCore v1.25.0
  [38540f10] CommonSolve v0.2.4
  [bbf7d656] CommonSubexpressions v0.3.1
  [34da2185] Compat v4.16.0
  [a33af91c] CompositionsBase v0.1.2
  [88cd18e8] ConsoleProgressMonitor v0.1.2
  [187b0558] ConstructionBase v1.5.8
  [9a962f9c] DataAPI v1.16.0
  [864edb3b] DataStructures v0.18.20
  [e2d170a0] DataValueInterfaces v1.0.0
  [163ba53b] DiffResults v1.1.0
  [b552c78f] DiffRules v1.15.1
  [a0c0ee7d] DifferentiationInterface v0.6.24
  [ffbed154] DocStringExtensions v0.9.3
  [4e289a0a] EnumX v1.0.4
  [e2ba6199] ExprTools v0.1.10
⌅ [6b7a57c9] Expronicon v0.8.5
  [9aa1b823] FastClosures v0.3.2
  [1a297f60] FillArrays v1.13.0
  [6a86dc24] FiniteDiff v2.26.2
  [f6369f11] ForwardDiff v0.10.38
  [069b7b12] FunctionWrappers v1.1.3
  [77dc65aa] FunctionWrappersWrappers v0.1.3
  [0c68f7d7] GPUArrays v11.1.0
  [46192b85] GPUArraysCore v0.2.0
  [7869d1d1] IRTools v0.4.14
  [3587e190] InverseFunctions v0.1.17
  [92d709cd] IrrationalConstants v0.2.2
  [82899510] IteratorInterfaceExtensions v1.0.0
  [692b3bcd] JLLWrappers v1.6.1
  [63c18a36] KernelAbstractions v0.9.29
  [5be7bae1] LBFGSB v0.4.1
  [929cbde3] LLVM v9.1.3
  [1d6d02ad] LeftChildRightSiblingTrees v0.2.0
  [d3d80556] LineSearches v7.3.0
  [2ab3a3ac] LogExpFunctions v0.3.28
  [e6f89c97] LoggingExtras v1.1.0
  [d8e11817] MLStyle v0.4.17
  [1914dd2f] MacroTools v0.5.13
  [e1d29d7a] Missings v1.2.0
  [d41bc354] NLSolversBase v7.8.3
  [77ba4419] NaNMath v1.0.2
  [429524aa] Optim v1.10.0
  [7f7a1694] Optimization v4.0.5
  [bca83a33] OptimizationBase v2.4.0
  [36348300] OptimizationOptimJL v0.4.1
  [bac558e1] OrderedCollections v1.7.0
  [90014a1f] PDMats v0.11.31
  [d96e819e] Parameters v0.12.3
  [85a6dd25] PositiveFactorizations v0.2.4
  [aea7be01] PrecompileTools v1.2.1
  [21216c6a] Preferences v1.4.3
  [33c8b6b6] ProgressLogging v0.1.4
  [92933f4c] ProgressMeter v1.10.2
  [c1ae055f] RealDot v0.1.0
  [3cdcf5f2] RecipesBase v1.3.4
  [731186ca] RecursiveArrayTools v3.27.4
  [189a3867] Reexport v1.2.2
  [ae029012] Requires v1.3.0
  [7e49a35a] RuntimeGeneratedFunctions v0.5.13
  [0bca4576] SciMLBase v2.65.0
  [c0aeaf25] SciMLOperators v0.3.12
  [53ae85a6] SciMLStructures v1.6.1
  [efcf1570] Setfield v1.1.1
  [a2af1166] SortingAlgorithms v1.2.1
  [9f842d2f] SparseConnectivityTracer v0.6.9
  [dc90abb0] SparseInverseSubset v0.1.2
  [0a514795] SparseMatrixColorings v0.4.10
  [276daf66] SpecialFunctions v2.4.0
  [90137ffa] StaticArrays v1.9.8
  [1e83bf80] StaticArraysCore v1.4.3
  [82ae8749] StatsAPI v1.7.0
  [2913bbd2] StatsBase v0.34.3
⌃ [09ab397b] StructArrays v0.6.21
  [2efcf032] SymbolicIndexingInterface v0.3.35
  [3783bdb8] TableTraits v1.0.1
  [bd369af6] Tables v1.12.0
  [5d786b92] TerminalLoggers v0.1.7
  [3a884ed6] UnPack v1.0.2
  [013be700] UnsafeAtomics v0.2.1
  [d80eeb9a] UnsafeAtomicsLLVM v0.2.1
  [e88e6eb3] Zygote v0.6.73
  [700de1a5] ZygoteRules v0.2.5
  [dad2f222] LLVMExtra_jll v0.0.34+0
  [81d17ec3] L_BFGS_B_jll v3.0.1+0
  [efe28fd5] OpenSpecFun_jll v0.5.5+0
  [0dad84c5] ArgTools v1.1.1
  [56f22d72] Artifacts
  [2a0f44e3] Base64
  [ade2ca70] Dates
  [8ba89e20] Distributed
  [f43a241f] Downloads v1.6.0
  [7b1f6079] FileWatching
  [9fa8497b] Future
  [b77e0a4c] InteractiveUtils
  [4af54fe1] LazyArtifacts
  [b27032c2] LibCURL v0.6.4
  [76f85450] LibGit2
  [8f399da3] Libdl
  [37e2e46d] LinearAlgebra
  [56ddb016] Logging
  [d6f4376e] Markdown
  [ca575930] NetworkOptions v1.2.0
  [44cfe95a] Pkg v1.10.0
  [de0858da] Printf
  [3fa0cd96] REPL
  [9a3f8284] Random
  [ea8e919c] SHA v0.7.0
  [9e88b42a] Serialization
  [6462fe0b] Sockets
  [2f01184e] SparseArrays v1.10.0
  [10745b16] Statistics v1.10.0
  [4607b0f0] SuiteSparse
  [fa267f1f] TOML v1.0.3
  [a4e569a6] Tar v1.10.0
  [cf7118a7] UUIDs
  [4ec0a83e] Unicode
  [e66e0078] CompilerSupportLibraries_jll v1.1.1+0
  [deac9b47] LibCURL_jll v8.4.0+0
  [e37daf67] LibGit2_jll v1.6.4+0
  [29816b5a] LibSSH2_jll v1.11.0+1
  [c8ffd9c3] MbedTLS_jll v2.28.2+1
  [14a3606d] MozillaCACerts_jll v2023.1.10
  [4536629a] OpenBLAS_jll v0.3.23+4
  [05823500] OpenLibm_jll v0.8.1+2
  [bea87d4a] SuiteSparse_jll v7.2.1+1
  [83775a58] Zlib_jll v1.2.13+1
  [8e850b90] libblastrampoline_jll v5.8.0+1
  [8e850ede] nghttp2_jll v1.52.0+1
  [3f19e933] p7zip_jll v17.4.0+2
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated -m`

• Output of versioninfo()

julia> versioninfo()
Julia Version 1.10.3
Commit 0b4590a5507 (2024-04-30 10:59 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: Linux (x86_64-linux-gnu)
  CPU: 24 × AMD Ryzen 9 3900XT 12-Core Processor
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, znver2)
Threads: 24 default, 0 interactive, 12 GC (on 24 virtual cores)

Additional context

Add any other context about the problem here.

[linjing-lab]

In large-scale optimization problems, x represents sequential parameter to be solved for the data (or 3D array), usually a long sequence, which is a sequence optimization problem that measures 3D problems. Therefore, using 3D arrays as parameters is not feasible in optimization problems.

[ChrisRackauckas]

That's not true. The optimizer can work perfectly fine on 3D, you just use the vec form. The issue is that the linear operator for the Hessian is not being properly constructed. That's not so hard to fix though. The actual issue is that you shouldn't be forming the Hessian here in the first place