Nested `ProductManifold`s

[lukas-weber]

The current implementation of the riemannian_gradient does not seem to handle nested ProductManifolds.

using Manifolds
using ManifoldDiff
using Zygote
using ADTypes

M = ProductManifold(ProductManifold(Euclidean(1), Euclidean(1)), Euclidean(1))
f(x) = sum(x)^2

ManifoldDiff.gradient(M, f, rand(M), RiemannianProjectionBackend(AutoZygote()))

The function `submanifold_components` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  submanifold_components(::ArrayPartition)
   @ ManifoldsBaseRecursiveArrayToolsExt /home/lweber/.julia/packages/ManifoldsBase/nqeOA/ext/ManifoldsBaseRecursiveArrayToolsExt/ProductManifoldRecursiveArrayToolsExt.jl:330
  submanifold_components(::ProductGroup, ::Identity{ProductOperation})
   @ Manifolds /home/lweber/.julia/packages/Manifolds/mH6Su/src/groups/product_group.jl:124
  submanifold_components(::GroupManifold{𝔽, MT, O}, ::Identity{O}) where {MT<:ProductManifold, 𝔽, O<:AbstractGroupOperation}
   @ Manifolds /home/lweber/.julia/packages/Manifolds/mH6Su/src/groups/product_group.jl:116
  ...

Stacktrace:
  [1] submanifold_components
    @ /home/lweber/.julia/packages/ManifoldsBase/nqeOA/src/ProductManifold.jl:1291 [inlined]
  [2] project!(M::ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, Y::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, p::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, X::Vector{Float64})
    @ ManifoldsBase /home/lweber/.julia/packages/ManifoldsBase/nqeOA/src/ProductManifold.jl:815
  [3] riemannian_gradient!(M::ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, X::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, p::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, Y::Vector{Float64}; embedding_metric::EuclideanMetric)
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:325
  [4] riemannian_gradient(M::ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, p::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, Y::Vector{Float64}; embedding_metric::EuclideanMetric)
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:314
  [5] riemannian_gradient(M::ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, p::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, Y::Vector{Float64})
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:307
  [6] map
    @ ./tuple.jl:406 [inlined]
  [7] riemannian_gradient(M::ProductManifold{ℝ, Tuple{ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, Vector{Float64}}}, X::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}})
    @ ManifoldsRecursiveArrayToolsExt /home/lweber/.julia/packages/Manifolds/mH6Su/ext/ManifoldsRecursiveArrayToolsExt/ProductManifoldRAT.jl:45
  [8] gradient(M::ProductManifold{ℝ, Tuple{ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, f::Function, p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, Vector{Float64}}}, backend::RiemannianProjectionBackend{AutoZygote}; kwargs::@Kwargs{})
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:207
  [9] gradient(M::ProductManifold{ℝ, Tuple{ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, f::Function, p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, Vector{Float64}}}, backend::RiemannianProjectionBackend{AutoZygote})
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:199
 [10] top-level scope
    @ REPL[117]:1

In particular, in ext/ManifoldsRecursiveArrayToolsExt/ProductManifoldRAT.jl:45, riemann_diff maps over M.manifolds, submanifold_components(p) and submanifold_components(X), but only the latter two are recursive.

[kellertuer]

I think the main challenge to me (and I will have to find time to track that further down) is that in your

[2] project!(M::ProductManifold{ℝ, Tuple{Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, Y::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, p::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}}, X::Vector{Float64})

with p being an ArrayPartition I am very surprised X is not – and as far as I see prior to that the Y.
So maybe the error is somewhere in [5] or [7], but I would even more suspect your f.

What you expect here to magically happen is that x (in R3?) is converted back and forth to RxRxR? Since both are Euclidean, the idea of using the embedding (which is the same as the manifold) for the gradient computation would here yield the identity, so we can not just guess that x is in R3 and not from M – and hence also have a problem to “convert the gradient back” – that is why X is still a vector.

I think this can be fixed, but in this specific example it might be a bit hard, since for Euclidean manifolds, like your M with a special form, we usually have that M == get_embedding(M) which would mean your x has to have the same type as your p (or rand(M)). Changing that might be a neat idea, but also probably breaking and I can not directly see how much work that would be.

edit: or for the specific example: The problem is that X is of wrong type, not that it is not recursive – a plan vector (or matrix) simply even semantically never has submanifold_components since that is simply to possible to magically guess where to split the vector/matrix. That is a function specifically tailored for “points/vectors with parts” like the ArrayPartition

[lukas-weber]

Ah, the minimal example is not chosen very well. What I actually need is ProductManifold{<:ProductManifold{<:Grassmanian,<:Grassmanian}}, <:Euclidean}. The same error appears there.

MethodError: no method matching submanifold_components(::Vector{Float64})
The function `submanifold_components` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  submanifold_components(::ArrayPartition)
   @ ManifoldsBaseRecursiveArrayToolsExt /home/lweber/.julia/packages/ManifoldsBase/nqeOA/ext/ManifoldsBaseRecursiveArrayToolsExt/ProductManifoldRecursiveArrayToolsExt.jl:330
  submanifold_components(::ProductGroup, ::Identity{ProductOperation})
   @ Manifolds /home/lweber/.julia/packages/Manifolds/mH6Su/src/groups/product_group.jl:124
  submanifold_components(::GroupManifold{𝔽, MT, O}, ::Identity{O}) where {MT<:ProductManifold, 𝔽, O<:AbstractGroupOperation}
   @ Manifolds /home/lweber/.julia/packages/Manifolds/mH6Su/src/groups/product_group.jl:116
  ...

Stacktrace:
  [1] submanifold_components
    @ /home/lweber/.julia/packages/ManifoldsBase/nqeOA/src/ProductManifold.jl:1291 [inlined]
  [2] project!(M::ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, Y::ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, p::ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, X::Vector{Float64})
    @ ManifoldsBase /home/lweber/.julia/packages/ManifoldsBase/nqeOA/src/ProductManifold.jl:815
  [3] riemannian_gradient!(M::ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, X::ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, p::ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Y::Vector{Float64}; embedding_metric::EuclideanMetric)
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:325
  [4] riemannian_gradient(M::ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, p::ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Y::Vector{Float64}; embedding_metric::EuclideanMetric)
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:314
  [5] riemannian_gradient(M::ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, p::ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Y::Vector{Float64})
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:307
  [6] map
    @ ./tuple.jl:406 [inlined]
  [7] riemannian_gradient(M::ProductManifold{ℝ, Tuple{ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Vector{Float64}}}, X::ArrayPartition{Float64, Tuple{Vector{Float64}, Vector{Float64}}})
    @ ManifoldsRecursiveArrayToolsExt /home/lweber/.julia/packages/Manifolds/mH6Su/ext/ManifoldsRecursiveArrayToolsExt/ProductManifoldRAT.jl:45
  [8] gradient(M::ProductManifold{ℝ, Tuple{ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, f::Function, p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Vector{Float64}}}, backend::RiemannianProjectionBackend{AutoZygote}; kwargs::@Kwargs{})
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:207
  [9] gradient(M::ProductManifold{ℝ, Tuple{ProductManifold{ℝ, Tuple{Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}, Grassmann{ManifoldsBase.TypeParameter{Tuple{2, 1}}, ℝ}}}, Euclidean{ManifoldsBase.TypeParameter{Tuple{1}}, ℝ}}}, f::Function, p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Vector{Float64}}}, backend::RiemannianProjectionBackend{AutoZygote})
    @ ManifoldDiff /home/lweber/.julia/packages/ManifoldDiff/vWWFF/src/riemannian_diff.jl:199
 [10] top-level scope
    @ REPL[121]:1

Regarding the cause, my bad, I did not look carefully enough. Actually, this is probably the fault of Zygote: it is supposed to spit out a tangent that is of the same type as p, which it usually does, but not for nested partitions it seems. So probably this is something for RecursiveArrayTools.jl to address in their Zygote extension!

[kellertuer]

Thanks! Yes and I should have asked for the more reasonable example as well :)

It might also be a bit our fault, I am not yet 100% sure, but yes, If they would return the same “shape” or “type” as their input, that would probably resolve that as well.

It might also be that we struggle a bit with the “hybrid” thing, an ArrayPartition that contains another ArrayPartition and a Vector

But from the signatures of the functions I see in that error, the same as my previous comment holds (and sorry if that was a bit of rambling), that I would expect X to be “of same type” as Y or here even as p as well.

p::ArrayPartition{Float64, Tuple{ArrayPartition{Float64, Tuple{Matrix{Float64}, Matrix{Float64}}}, Vector{Float64}}}

looks a bit frightening, but is nicely structured in [8] or [9] and X (in [7] for example) should for best of cases have a similar “structure”.

I do see an alternative, but that is a bit of a leeway.
our embed and project´ do a bit more than the mathematical forms (Euclidean would stay as is, Grassmann is embedded in the space of matrices) it also cares for the “right types”, though for now I can only explain that really quickly with fixed rank: Fixed rank matrices are stored as their (shortened, to nonzero singular values) SVD. So project` would not only “cut the rank down” and only keep the largest singular values but also “decompose” into the SVD representation.

What we could do, is, define an EmbeddedManifold(YourManifold, Euclidean(N)), where N is the sum of all entries of the two matrices and your Euclidean space,
and then define a project and embed that would turn such a vector into an ArrayPartition, so that Zygote could work on the N-vector and riemannian_gradient would automatically care for the “right structure” since it uses embed/project.

I hope that idea came across?

But before I (or someone else) dive(s) down such a rabbit hole, getting Zygote to return a gradient X as same type as p would be the much easier and best solution.

[mateuszbaran]

Until the issue with Zygote/RecursiveArrayTools.jl is fixed, you may also consider using a ProductManifold of three manifolds instead of a nested one.