Optimizing over 3d array: adjoint not defined for Array{Float64, 3}

I’m trying to optimize over a 3d array. It seems like it doesn’t want to handle a 3d array, but I’m not sure why. Is there a solution for this?

using Optimization, Optim, OptimizationOptimJL, Zygote, DifferentiationInterface
let k = 5
    f(Ψ, _) = sum((@view Ψ[:,:,i]) * (@view Ψ[:,:,i]) for i in 1:k)
    @assert all(isfinite, f(rand(5,5,k), nothing))
    optfunc = OptimizationFunction(f, Optimization.AutoZygote())
    problem = OptimizationProblem(optfunc, rand(5,5,k))
    solve(problem, BFGS())
end
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
1 Like

Perhaps you could try using a vector of matrices instead of a 3d array?

You can also just use Optim.jl directly, without adding so many layers of interface, since one of them seems to be the problem.

f returns a matrix, so I’m not sure what you want to minimise, perhaps the sum? That can be made arbitrarily large negative.

fjar(Ψ) = sum((@view Ψ[:,:,i]) * (@view Ψ[:,:,i]) for i in axes(Ψ,3))
function fmine(Ψ)  # faster version with no slices
    # this implements  @tensor out[a,c] := Ψ[a,b,i] * Ψ[b,c,i]
    A = reshape(Ψ, size(Ψ,1), :)  # A[a, (b,i)]
    B = reshape(permutedims(Ψ, (1,3,2)), :, size(Ψ,2))   # B[(b,i), c]
    A * B
end

using Optim, ForwardDiff
res = let k = 5, f = sum ∘ fjar
    x0 = rand(5,5,k)
    @show f(x0) # 157.8, on one run
    res = optimize(f, x0, LBFGS(), autodiff = :forward)
end
res.minimum  # -9.3e68, i.e. it goes to negative infinity

using Zygote
let k = 5, f = f = sum ∘ fjar  # same thing, with Zygote
    x0 = rand(5,5,k)
    grad!(dΨ, Ψ) = copyto!(dΨ, Zygote.gradient(f, Ψ)[1])
    res = optimize(f, grad!, x0, LBFGS())
end
1 Like