Ambiguous type inference for optimization over a PowerManifold

[yuwenchen95]

I want to optimize a quadratic cost <x,Cx> w.r.t. a matrix x of dimension \mathbb{R}^{n*d x rnk} where each subarray x[d*(i-1)+1:i*d,:] is constrained to a Stiefel manifold x[d*(i-1)+1:i*d,:]*x[d*(i-1)+1:i*d,:]' = I_d. In my implementation, I use a 3d arrays of dimension (rnk,d,n) to represent a valid x. Here is my code with recording ang debugging:

d = 3
n = 100
dim = 3*100
C = sprandn(dim, dim, 0.1)
C = 0.5*(C+C');     #symmetrize if not symmetric
C = (-0.25).*C;

rnk = Int(ceil(sqrt(2*dim)))
M = PowerManifold(Stiefel(rnk,d),n) 
function F(M,x) 
    xred = reshape(x,rnk,d*n)
    return tr(xred*C*xred')
end

struct grad!{TD,TTMP}
    C::TD
    tmp::TTMP
end

function (gradF!::grad!)(M,y,x)
    xred = reshape(x,rnk,d*n)
    yred = reshape(y,rnk,d*n)
    mul!(yred,xred,gradF!.C,2.,0.)
    Threads.@threads for i = 1:n
        id = Threads.threadid()

        cur_x = @view x[:,:,i]
        cur_y = @view y[:,:,i]

        dotxy = @view gradF!.tmp[:,:,id]
        mul!(dotxy,cur_x',cur_y)
        @. dotxy = (dotxy + dotxy')/2.
        cur_y .-= cur_x*dotxy
    end

    return y
end

gradF2! = grad!(C,zeros(d,d,Threads.nthreads()))

x0 = rand(rnk,d,n)
for i = 1:n
    cur_x = @view x0[:,:,i]
    (Ui,Si,VTi) = LAPACK.gesdd!('S',cur_x)
    mul!(cur_x,Ui,VTi)
end

dg=[:Iteration, " | ", :Change, " | ", :Cost, "\n", :Stop]
rd=[:x, :Change, :Cost]
ro=true
st = ArmijoLinesearch(1.0, ExponentialRetraction(), 0.5, 0.001)
sc = StopAfterIteration(50) | StopWhenGradientNormLess(1e-10)

m = deepcopy(x0)
@time rgd = gradient_descent!(
    M, F, gradF2!, m; evaluation=MutatingEvaluation(), stepsize=st, stopping_criterion=sc,
    record = rd, return_options=true, debug= dg
)
rgd_cost = get_record(rgd, :Iteration, :Cost)

However, it encountered the following error:

julia> include("test_manifolds.jl")
Initial  |  | F(x): -4.310718
# 1      | ERROR: LoadError: MethodError: log!(::AbstractManifold, ::Matrix{Float64}, ::SubArray{Float64, 2, Array{Float64, 3}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Int64}, true}, ::SubArray{Float64, 2, Array{Float64, 3}, Tuple{Base.Slice{Base.OneTo{Int64}}, Base.Slice{Base.OneTo{Int64}}, Int64}, true}) is ambiguous. Candidates:
  log!(::CholeskySpace, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\CholeskySpace.jl:147
  log!(M::Hyperbolic, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\HyperbolicHyperboloid.jl:373
  log!(::TangentSpaceAtPoint{𝔽} where 𝔽  , X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\VectorBundle.jl:718
  log!(M::ValidationManifold, X, p, q; kwargs...) in ManifoldsBase at C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\ValidationManifold.jl:297
  log!(::Manifolds.GeneralUnitaryMatrices{4, ℝ}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GeneralUnitaryMatrices.jl:569
  log!(M::Manifolds.GeneralUnitaryMatrices{3, ℝ}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GeneralUnitaryMatrices.jl:556
  log!(M::Manifolds.GeneralUnitaryMatrices{2, ℝ}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GeneralUnitaryMatrices.jl:550
  log!(::Manifolds.GeneralUnitaryMatrices{n, ℝ}, X, p, q) where n in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GeneralUnitaryMatrices.jl:545
  log!(::Manifolds.GeneralUnitaryMatrices{n, 𝔽} , X, p, q) where {n, 𝔽 } in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GeneralUnitaryMatrices.jl:595
  log!(::GeneralLinear{1}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\groups\general_linear.jl:217
  log!(G::GeneralLinear{n, 𝔽} , X, p, q) where {n, 𝔽 } in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\groups\general_linear.jl:198
  log!(M::ConnectionManifold{𝔽, <:AbstractDecoratorManifold{𝔽}, <:AbstractCartanSchoutenConnection}  , Y, p, q) where 𝔽
 in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\groups\connections.jl:102
  log!(::Grassmann{n, k}, X, p, q) where {n, k} in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GrassmannStiefel.jl:174
  log!(::SymmetricPositiveSemidefiniteFixedRank, Z, q, p) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\SymmetricPositiveSemidefiniteFixedRank.jl:175
  log!(M::AbstractSphere, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\Sphere.jl:312
  log!(::ManifoldsBase.DefaultManifold, Y, p, q) in ManifoldsBase at C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\DefaultManifold.jl:94
  log!(M::HeisenbergGroup{n}, X, p, q) where n in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\groups\heisenberg.jl:261
  log!(::Euclidean, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\Euclidean.jl:427
  log!(::ProbabilitySimplex, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\ProbabilitySimplex.jl:240
  log!(::SymmetricPositiveDefinite{N}, X, p, q) where N in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\SymmetricPositiveDefiniteLinearAffine.jl:350
  log!(M::AbstractProjectiveSpace, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\ProjectiveSpace.jl:313
  log!(::MetricManifold{ℝ, SymmetricPositiveDefinite{N}, LogCholeskyMetric}, X, p, q) where N in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\SymmetricPositiveDefiniteLogCholesky.jl:119
  log!(M::MetricManifold{ℝ, <:SymmetricPositiveDefinite, <:GeneralizedBuresWassersteinMetric}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\SymmetricPositiveDefiniteGeneralizedBuresWasserstein.jl:133
  log!(::MetricManifold{ℝ, <:SymmetricPositiveDefinite, BuresWassersteinMetric}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\SymmetricPositiveDefiniteBuresWasserstein.jl:130
  log!(M::MetricManifold{ℝ, Stiefel{n, k, ℝ}, CanonicalMetric}, X, p, q; maxiter, tolerance) where {n, k} in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\StiefelCanonicalMetric.jl:142
  log!(M::GeneralizedGrassmann{n, k}, X, p, q) where {n, k} in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\GeneralizedGrassmann.jl:238
  log!(M::ProductManifold{ℝ, Tuple{GroupManifold{ℝ, Euclidean{Tuple{N}, ℝ}, AdditionOperation}, SpecialOrthogonal{N}}} where N, X::AbstractMatrix, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\groups\special_euclidean.jl:533
  log!(M::ProductManifold, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\ProductManifold.jl:802
  log!(M::AbstractDecoratorManifold, X, p, q) in ManifoldsBase at C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:274
  log!(M::EssentialManifold, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\EssentialManifold.jl:168
  log!(M::AbstractPowerManifold{𝔽, <:AbstractManifold{𝔽}, NestedReplacingPowerRepresentation} where 𝔽   , X, p, q) in ManifoldsBase at C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\PowerManifold.jl:812
  log!(M::AbstractPowerManifold, X, p, q) in ManifoldsBase at C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\PowerManifold.jl:800
  log!(::PositiveNumbers, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\PositiveNumbers.jl:185
  log!(::Circle{ℝ}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\Circle.jl:272
  log!(M::Circle{ℂ}, X, p, q) in Manifolds at C:\Users\lalinc4817\.julia\packages\Manifolds\Dypjk\src\manifolds\Circle.jl:273
Stacktrace:
  [1] log!
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:308 [inlined]
  [2] log! (repeats 4 times)
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:282 [inlined]
  [3] log!
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:275 [inlined]
  [4] log
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\exp_log_geo.jl:143 [inlined]
  [5] log
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:308 [inlined]
  [6] log (repeats 4 times)
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:282 [inlined]
  [7] log
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\nested_trait.jl:275 [inlined]
  [8] distance
    @ C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\ManifoldsBase.jl:285 [inlined]
  [9] distance(M::PowerManifold{ℝ, Stiefel{25, 3, ℝ}, Tuple{100}, ArrayPowerRepresentation}, p::Array{Float64, 3}, q::Array{Float64, 3})
    @ ManifoldsBase C:\Users\lalinc4817\.julia\packages\ManifoldsBase\YRaAP\src\PowerManifold.jl:380
 [10] (::DebugChange)(p::GradientProblem{MutatingEvaluation, PowerManifold{ℝ, Stiefel{25, 3, ℝ}, Tuple{100}, ArrayPowerRepresentation}, typeof(F), grad!{SparseMatrixCSC{Float64, Int64}, Array{Float64, 3}}}, o::GradientDescentOptions{Array{Float64, 3}, Array{Float64, 3}, StopWhenAny{Tuple{StopAfterIteration, StopWhenGradientNormLess}}, ArmijoLinesearch{ExponentialRetraction}, ExponentialRetraction}, i::Int64)
    @ Manopt C:\Users\lalinc4817\.julia\packages\Manopt\83Ybx\src\plans\debug_options.jl:148
 [11] (::DebugGroup)(p::GradientProblem{MutatingEvaluation, PowerManifold{ℝ, Stiefel{25, 3, ℝ}, Tuple{100}, ArrayPowerRepresentation}, typeof(F), grad!{SparseMatrixCSC{Float64, Int64}, Array{Float64, 3}}}, o::GradientDescentOptions{Array{Float64, 3}, Array{Float64, 3}, StopWhenAny{Tuple{StopAfterIteration, StopWhenGradientNormLess}}, ArmijoLinesearch{ExponentialRetraction}, ExponentialRetraction}, i::Int64)
    @ Manopt C:\Users\lalinc4817\.julia\packages\Manopt\83Ybx\src\plans\debug_options.jl:88
 [12] step_solver!(p::GradientProblem{MutatingEvaluation, PowerManifold{ℝ, Stiefel{25, 3, ℝ}, Tuple{100}, ArrayPowerRepresentation}, typeof(F), grad!{SparseMatrixCSC{Float64, Int64}, Array{Float64, 3}}}, o::DebugOptions{GradientDescentOptions{Array{Float64, 3}, Array{Float64, 3}, StopWhenAny{Tuple{StopAfterIteration, StopWhenGradientNormLess}}, ArmijoLinesearch{ExponentialRetraction}, ExponentialRetraction}}, i::Int64)
    @ Manopt C:\Users\lalinc4817\.julia\packages\Manopt\83Ybx\src\solvers\debug_solver.jl:24
 [13] step_solver!
    @ C:\Users\lalinc4817\.julia\packages\Manopt\83Ybx\src\solvers\record_solver.jl:19 [inlined]
 [14] solve(p::GradientProblem{MutatingEvaluation, PowerManifold{ℝ, Stiefel{25, 3, ℝ}, Tuple{100}, ArrayPowerRepresentation}, typeof(F), grad!{SparseMatrixCSC{Float64, Int64}, Array{Float64, 3}}}, o::RecordOptions{DebugOptions{GradientDescentOptions{Array{Float64, 3}, Array{Float64, 3}, StopWhenAny{Tuple{StopAfterIteration, StopWhenGradientNormLess}}, ArmijoLinesearch{ExponentialRetraction}, ExponentialRetraction}}, NamedTuple{(:Iteration,), Tuple{RecordGroup}}})
    @ Manopt C:\Users\lalinc4817\.julia\packages\Manopt\83Ybx\src\solvers\solver.jl:88
 [15] gradient_descent!(M::PowerManifold{ℝ, Stiefel{25, 3, ℝ}, Tuple{100}, ArrayPowerRepresentation}, F::typeof(F), gradF::grad!{SparseMatrixCSC{Float64, Int64}, Array{Float64, 3}}, x::Array{Float64, 3}; stepsize::ArmijoLinesearch{ExponentialRetraction}, retraction_method::ExponentialRetraction, stopping_criterion::StopWhenAny{Tuple{StopAfterIteration, StopWhenGradientNormLess}}, debug::Vector{Any}, direction::IdentityUpdateRule, evaluation::MutatingEvaluation, return_options::Bool, kwargs::Base.Pairs{Symbol, Vector{Symbol}, Tuple{Symbol}, NamedTuple{(:record,), Tuple{Vector{Symbol}}}})
    @ Manopt C:\Users\lalinc4817\.julia\packages\Manopt\83Ybx\src\solvers\gradient_descent.jl:88
 [16] top-level scope
    @ .\timing.jl:263
 [17] include(fname::String)
    @ Base.MainInclude .\client.jl:476
 [18] top-level scope
    @ REPL[18]:1
in expression starting at E:\code\paper_code\paper_admm_bm-main\test_manifolds.jl:131

It seems to be a a type inference issue related to traits mensioned in this page of ManifoldBase.jl. How could I solve it?

[kellertuer]

Thanks for spotting this, I already narrowed it down to an M(n)WE at

JuliaManifolds/Manifolds.jl#537

since that is an issue with Manifolds.jl and we will fix that soon I hope.

[kellertuer]

Ah, With the help of @mateuszbaran we found that the main problem is, that
in order to get the :Change debug, Manopt.jl relies on either distance being implemented and if that is not implemented a log (and the norm on vectors) being implemented.

Neither of them exists on the Stiefel manifold, at least not for the default Euclidean metric – but we have something like that soon based on two papers from the last two years, see JuliaManifolds/Manifolds.jl#535), on the Canonical Metric there is an approximate solution to the log, but for that you have to use the https://manoptjl.org/stable/plans/debug/#Manopt.DebugEntryChange instead (which is very general and not that easy) in order to compute the distance with your own function.

edit: Or long story short: Neither distance nor log are implemented so that is where the error comes from. Both are not so easy, since approxiations to these just were published last year.

[kellertuer]

Addon number 2, you could do the following

Define (after importing both)

ManifoldsBase.default_inverse_retraction_method(M::Stiefel) = PolarInverseRetraction()
ManifoldsBase.default_inverse_retraction_method(M::PowerManifold) = ManifoldsBase.default_inverse_retraction_method(M.manifold)

then

distance(M, x0, x0, ManifoldsBase.default_inverse_retraction_method(M))

does work as a rough estimate (which should be enough) and that you could use in the DebugEntryChange that I linked above.

[kellertuer]

The problem with DebugEntryChange is that one has to provide a manual distance function (if the default is not used), but right, maybe JuliaManifolds/ManifoldsBase.jl#125 helps here. For the approximation of the Hessian: Since that also uses the `Default_retraction_method) the PR mentioned above might also help there.

[kellertuer]

With that PR merged, a few things on PowerManifold might work more smoothly for Stiefel (since now a proper default retraction / inverse retraction / vector transport is chosen that exists); this is used also in the approximate Hessian, so that might work now, for example (using ManifoldsBase 0.13.20).

[yuwenchen95]

I have updated both Manopt and ManifoldsBase packages and it works now. Thank you so much for your patient explanation!

(A minor issue is that Stiefel manifold doesn't have a support for ParallelTransport, which is the default setting for the vector_transport_method, and it would still fall into the case of ambiguous reference if I didn't set it specifically, like ProjectionTransport. It would be nice if we could return a warning of it)

[kellertuer]

We can add that default in the next release as well, no problem.

edit: Will be part of JuliaManifolds/Manifolds.jl#538

[kellertuer]

with the linked one merged, the usual fallbacks for Stiefel are introduced, to that default_vectortransport_method(::Stiefel) returns an existing one – with the check on Power that we already have that should work as well, so most things should work fine now (the change debug not yet, but soon as well).