feat: GPU-assisted exp! for PowerManifold(SymmetricPositiveDefinite)

Serial batch implementation: each SPD exp uses CPU eigendecomposition
(no batched GPU eigendecomp available in CUDA.jl at time of writing).

Formula: exp_p(X) = p^{1/2} · expm(p^{-1/2} · X · p^{-1/2}) · p^{1/2}
Output is symmetrized (0.5*(result + result')) to ensure exact symmetry
within Manifolds.jl's is_point tolerance.

TODO: replace with cusolverDnSsyevjBatched when exposed in CUDA.jl.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: zazabap

ext/ManifoldsGPUCUDAExt/ManifoldsGPUCUDAExt.jl

@@ -12,6 +12,7 @@ include("Grassmann.jl")
 include("Sphere.jl")
 include("Euclidean.jl")
 include("UnitaryMatrices.jl")
+include("SPD.jl")
 
 
 end

ext/ManifoldsGPUCUDAExt/SPD.jl

@@ -0,0 +1,50 @@
+using LinearAlgebra
+
+"""
+    exp!(M::PowerManifold{ℝ, <:SymmetricPositiveDefinite, ...}, q, p, X)
+
+GPU-assisted geodesic exponential on batched SPD matrices.
+
+Strategy: serial loop over batch — each (n,n) SPD exp uses CPU eigendecomposition
+via eigen(Symmetric(...)). Slices are transferred CPU↔GPU per batch element.
+
+TODO: Replace with true GPU-batched implementation when CUDA.CUSOLVER exposes
+syevjBatched (batched symmetric eigendecomposition).
+
+Formula: exp_p(X) = p^{1/2} · expm(p^{-1/2} · X · p^{-1/2}) · p^{1/2}
+"""
+function ManifoldsBase.exp!(
+    ::PowerManifold{ℝ, <:SymmetricPositiveDefinite, <:Tuple, ArrayPowerRepresentation},
+    q::CuArray{T, 3},
+    p::CuArray{T, 3},
+    X::CuArray{T, 3},
+) where {T <: Real}
+    batch = size(q, 3)
+    for i in 1:batch
+        # Transfer to CPU for eigendecomposition (no batched GPU alternative yet)
+        p_i = Array(@view p[:, :, i])
+        X_i = Array(@view X[:, :, i])
+
+        # Eigendecomposition of p_i (symmetric PD): p = V * diag(λ) * V'
+        λ, V = eigen(Symmetric(p_i))
+        sqrt_λ    = sqrt.(λ)
+        invsqrt_λ = inv.(sqrt_λ)
+        p_sqrt    = V * Diagonal(sqrt_λ)    * V'
+        p_invsqrt = V * Diagonal(invsqrt_λ) * V'
+
+        # Symmetric inner matrix: S = p^{-1/2} · X · p^{-1/2}
+        S = Symmetric(p_invsqrt * X_i * p_invsqrt)
+
+        # Matrix exponential of symmetric S via eigendecomposition
+        μ, W  = eigen(S)
+        expS  = W * Diagonal(exp.(μ)) * W'
+
+        # Geodesic: q = p^{1/2} · exp(S) · p^{1/2}
+        # Symmetrize to cancel floating-point rounding: product of symmetric
+        # matrices is only approximately symmetric in finite precision.
+        result = p_sqrt * expS * p_sqrt
+        @view(q[:, :, i]) .= CuArray(T.(0.5 .* (result .+ result')))
+    end
+
+    return q
+end

test/cuda_tests.jl

@@ -327,4 +327,38 @@ using CUDA
         @test is_point(MP, Array(Y_cu))
         @test isapprox(MP, p, Array(Y_cu), Y_cpu; atol = 2.0f-4, rtol = 2.0f-4)
     end
+
+    @testset "SPD exp Float64" begin
+        Random.seed!(90)
+        M = SymmetricPositiveDefinite(4)
+        MP = PowerManifold(M, 16)
+
+        p = rand(MP)
+        X = 0.2 .* rand(MP; vector_at = p)
+        Y_cpu = exp(MP, p, X)
+
+        p_cu = CuArray(p)
+        X_cu = CuArray(X)
+        Y_cu = exp(MP, p_cu, X_cu)
+
+        @test is_point(MP, Array(Y_cu))
+        @test isapprox(MP, p, Array(Y_cu), Y_cpu; atol = 2.0e-12, rtol = 2.0e-12)
+    end
+
+    @testset "SPD exp Float32" begin
+        Random.seed!(91)
+        M = SymmetricPositiveDefinite(4)
+        MP = PowerManifold(M, 16)
+
+        p = Float32.(rand(MP))
+        X = Float32(0.2) .* Float32.(rand(MP; vector_at = p))
+        Y_cpu = exp(MP, p, X)
+
+        p_cu = CuArray(p)
+        X_cu = CuArray(X)
+        Y_cu = exp(MP, p_cu, X_cu)
+
+        @test is_point(MP, Array(Y_cu))
+        @test isapprox(MP, p, Array(Y_cu), Y_cpu; atol = 2.0f-5, rtol = 2.0f-5)
+    end
 end

test/jlarray_tests.jl

@@ -193,3 +193,22 @@ end
     @test is_point(MP, Array(Y_jl))
     @test isapprox(MP, p, Array(Y_jl), Y_cpu; atol = 2.0e-12, rtol = 2.0e-12)
 end
+
+@testset "JLArray: SPD exp Float64" begin
+    M = SymmetricPositiveDefinite(4)
+    MP = PowerManifold(M, 8)
+
+    Random.seed!(90)
+    p = rand(MP)
+    X = 0.2 .* rand(MP; vector_at = p)
+    Y_cpu = exp(MP, p, X)
+
+    p_jl = JLArray(p)
+    X_jl = JLArray(X)
+    Y_jl = exp(MP, p_jl, X_jl)
+
+    # Note: no is_point check here — the CPU Manifolds.jl exp for SPD produces
+    # results with ~O(eps) asymmetry that fails Manifolds.jl's strict symmetry
+    # tolerance. The GPU implementation (SPD.jl) symmetrizes explicitly.
+    @test isapprox(MP, p, Array(Y_jl), Y_cpu; atol = 2.0e-12, rtol = 2.0e-12)
+end