From e1daf804236582c7a98774050f3d69af2b5e84f9 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Mon, 17 Mar 2025 10:46:17 +0100
Subject: [PATCH 1/8] feat: implement retraction selection mechanism

---
 src/natural_manifolds.jl        | 51 +++++++++++++++++++++++++++------
 src/natural_manifolds/normal.jl |  4 ++-
 2 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl
index 61fa791..5cd0642 100644
--- a/src/natural_manifolds.jl
+++ b/src/natural_manifolds.jl
@@ -8,20 +8,30 @@ import ExponentialFamily: exponential_family_typetag
 The manifold for the natural parameters of the distribution of type `T` with dimensions `dims`.
 An internal structure, use `get_natural_manifold` to create an instance of a manifold for the natural parameters of distribution of type `T`.
 """
-struct NaturalParametersManifold{𝔽,T,D,M,C} <: AbstractDecoratorManifold{𝔽}
+struct NaturalParametersManifold{𝔽,T,D,M,C,R} <: AbstractDecoratorManifold{𝔽}
     dims::D
     base::M
     conditioner::C
+    retraction::R
 end
 
 getdims(M::NaturalParametersManifold) = M.dims
 getbase(M::NaturalParametersManifold) = M.base
 getconditioner(M::NaturalParametersManifold) = M.conditioner
+getretraction(M::NaturalParametersManifold) = M.retraction
 
 # The `NaturalParametersManifold` simply adds extra properties to the `base` and 
 # acts as a "decorator"
-@inline ManifoldsBase.active_traits(f::F, ::NaturalParametersManifold, ::Any...) where {F} =
-    ManifoldsBase.IsExplicitDecorator()
+@inline function ManifoldsBase.active_traits(f::F, ::NaturalParametersManifold, args...) where {F}
+    # Don't delegate retraction-related methods
+    if f in (ManifoldsBase.retract, ManifoldsBase.retract!, 
+             ManifoldsBase.retract_fused, ManifoldsBase.retract_fused!)
+        return ManifoldsBase.EmptyTrait()
+    else
+        return ManifoldsBase.IsExplicitDecorator()
+    end
+end
+
 @inline ManifoldsBase.decorated_manifold(M::NaturalParametersManifold) = M.base
 
 function ExponentialFamily.exponential_family_typetag(
@@ -31,9 +41,9 @@ function ExponentialFamily.exponential_family_typetag(
 end
 
 function NaturalParametersManifold(
-    ::Type{T}, dims::D, base::M, conditioner::C=nothing
-) where {T,𝔽,D,M<:AbstractManifold{𝔽},C}
-    return NaturalParametersManifold{𝔽,T,D,M,C}(dims, base, conditioner)
+    ::Type{T}, dims::D, base::M, conditioner::C=nothing, retraction::R=nothing
+) where {T,𝔽,D,M<:AbstractManifold{𝔽},C,R}
+    return NaturalParametersManifold{𝔽,T,D,M,C,R}(dims, base, conditioner, retraction)
 end
 
 """
@@ -52,9 +62,9 @@ julia> ExponentialFamilyManifolds.get_natural_manifold(MvNormalMeanCovariance, (
 true
 ```
 """
-function get_natural_manifold(::Type{T}, dims, conditioner=nothing) where {T}
+function get_natural_manifold(::Type{T}, dims, conditioner=nothing, retraction = nothing) where {T}
     return NaturalParametersManifold(
-        T, dims, get_natural_manifold_base(T, dims, conditioner), conditioner
+        T, dims, get_natural_manifold_base(T, dims, conditioner), conditioner, retraction
     )
 end
 
@@ -88,3 +98,28 @@ function Base.convert(
         exponential_family_typetag(M), p, getconditioner(M), nothing
     )
 end
+
+struct ChartNOrderRetraction{Order, E} <: AbstractRetractionMethod
+    extra::E
+end
+
+function ChartNOrderRetraction{O}() where {O}
+    return ChartNOrderRetraction{O, Nothing}(nothing)
+end
+
+const FirstOrderRetraction = ChartNOrderRetraction{1}
+
+ManifoldsBase.default_retraction_method(::NaturalParametersManifold{𝔽,TD,D,M,C,Nothing}, ::Type{T}) where {𝔽, T,TD, D, M, C} = FirstOrderRetraction()
+
+ManifoldsBase.default_retraction_method(M::NaturalParametersManifold{𝔽,TD,D,BM,C,R}, ::Type{T}) where {𝔽, T,TD, D, BM, C, R} = getretraction(M)
+
+function ManifoldsBase.retract_fused!(::NaturalParametersManifold, q, p, X, t::Number, method::FirstOrderRetraction)
+    q .= p .+ t .* X
+    return q
+end
+
+function ManifoldsBase.retract!(M::NaturalParametersManifold, q, p, X, method::FirstOrderRetraction)
+    return ManifoldsBase.retract_fused!(M, q, p, X, one(eltype(X)), method)
+end
+
+
diff --git a/src/natural_manifolds/normal.jl b/src/natural_manifolds/normal.jl
index 0a91d9a..f7e2407 100644
--- a/src/natural_manifolds/normal.jl
+++ b/src/natural_manifolds/normal.jl
@@ -65,4 +65,6 @@ function partition_point(
 )
     k = first(dims)
     return ArrayPartition(view(p, 1:k), view(p, k+1:k+1))
-end
\ No newline at end of file
+end
+
+ManifoldsBase.default_retraction_method(::NaturalParametersManifold{𝔽, MvNormalMeanCovariance, D, M, C, Nothing}, ::Type{T}) where {𝔽, T, D, M, C} = ManifoldsBase.ExponentialRetraction()
\ No newline at end of file

From 582e0bdab08b825f377e9da14d7f6cf04e3a1ad8 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Mon, 17 Mar 2025 14:57:07 +0100
Subject: [PATCH 2/8] feat: draft Fisher Rao metric

---
 Project.toml                                  |  6 ++
 src/natural_manifolds.jl                      | 86 +++++++++++++++++--
 src/natural_manifolds/normal.jl               |  6 +-
 .../natural_manifolds_setuptests.jl           | 10 ++-
 test/natural_manifolds/normal_tests.jl        | 10 ++-
 5 files changed, 105 insertions(+), 13 deletions(-)

diff --git a/Project.toml b/Project.toml
index b584cc3..37e8b5a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "2.0.0"
 [deps]
 BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e"
 ExponentialFamily = "62312e5e-252a-4322-ace9-a5f4bf9b357b"
+FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb"
@@ -13,10 +14,15 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 
+[weakdeps]
+ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd"
+
 [compat]
 BayesBase = "1.3"
 ExponentialFamily = "2.0.0"
+FastCholesky = "1.3.1"
 LinearAlgebra = "1.10"
+ManifoldDiff = "0.4.2"
 Manifolds = "0.10"
 ManifoldsBase = "1"
 Random = "1.10"
diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl
index b02d44f..037bf0b 100644
--- a/src/natural_manifolds.jl
+++ b/src/natural_manifolds.jl
@@ -1,18 +1,43 @@
 using ManifoldsBase, Manifolds, Static, RecursiveArrayTools, Random, ExponentialFamily
+using FastCholesky
 
 import ExponentialFamily: exponential_family_typetag
 
+"""
+    FisherInformationMetric <: RiemannianMetric
+
+Specifier that we need to use the Fisher information metric.
+"""
+struct FisherInformationMetric <: RiemannianMetric end
+
+"""
+    BaseMetric <: RiemannianMetric
+
+Specifier that we need to use the metric from the base manifold.
+"""
+struct BaseMetric <: RiemannianMetric end
+
+"""
+    getdefaultmetric(::Type{T}) where {T}
+
+Returns the default metric for the distribution of type `T`.
+"""
+function getdefaultmetric(::Type{T}) where {T}
+    return FisherInformationMetric()
+end
+
 """
     NaturalParametersManifold(::Type{T}, dims, base, conditioner)
 
 The manifold for the natural parameters of the distribution of type `T` with dimensions `dims`.
 An internal structure, use `get_natural_manifold` to create an instance of a manifold for the natural parameters of distribution of type `T`.
 """
-struct NaturalParametersManifold{𝔽,T,D,M,C,R} <: AbstractDecoratorManifold{𝔽}
+struct NaturalParametersManifold{𝔽,T,D,M,C,R,MT} <: AbstractDecoratorManifold{𝔽}
     dims::D
     base::M
     conditioner::C
     retraction::R
+    metric::MT
 end
 
 getdims(M::NaturalParametersManifold) = M.dims
@@ -22,15 +47,18 @@ getretraction(M::NaturalParametersManifold) = M.retraction
 
 # The `NaturalParametersManifold` simply adds extra properties to the `base` and 
 # acts as a "decorator"
-@inline function ManifoldsBase.active_traits(
-    f::F, ::NaturalParametersManifold, args...
-) where {F}
-    # Don't delegate retraction-related methods
+function select_skip_methods(::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,BaseMetric}) where {F,𝔽,T,D,MB,C,R}
+    return ManifoldsBase.IsExplicitDecorator()
+end
+
+function select_skip_methods(f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}) where {F,𝔽,T,D,MB,C,R}
     if f in (
         ManifoldsBase.retract,
         ManifoldsBase.retract!,
         ManifoldsBase.retract_fused,
         ManifoldsBase.retract_fused!,
+        Manifolds.local_metric,
+        Manifolds.local_metric_jacobian,
     )
         return ManifoldsBase.EmptyTrait()
     else
@@ -38,6 +66,12 @@ getretraction(M::NaturalParametersManifold) = M.retraction
     end
 end
 
+@inline function ManifoldsBase.active_traits(
+    f::F, M::NaturalParametersManifold, args...
+) where {F}
+    return select_skip_methods(f, M)
+end
+
 @inline ManifoldsBase.decorated_manifold(M::NaturalParametersManifold) = M.base
 
 function ExponentialFamily.exponential_family_typetag(
@@ -47,9 +81,16 @@ function ExponentialFamily.exponential_family_typetag(
 end
 
 function NaturalParametersManifold(
-    ::Type{T}, dims::D, base::M, conditioner::C=nothing, retraction::R=nothing
-) where {T,𝔽,D,M<:AbstractManifold{𝔽},C,R}
-    return NaturalParametersManifold{𝔽,T,D,M,C,R}(dims, base, conditioner, retraction)
+    ::Type{T},
+    dims::D,
+    base::M,
+    conditioner::C=nothing,
+    retraction::R=nothing,
+    metric::MT=getdefaultmetric(T),
+) where {T,𝔽,D,M<:AbstractManifold{𝔽},C,R,MT}
+    return NaturalParametersManifold{𝔽,T,D,M,C,R,MT}(
+        dims, base, conditioner, retraction, metric
+    )
 end
 
 """
@@ -118,7 +159,7 @@ end
 const FirstOrderRetraction = ChartNOrderRetraction{1}
 
 function ManifoldsBase.default_retraction_method(
-    ::NaturalParametersManifold{𝔽,TD,D,M,C,Nothing}, ::Type{T}
+    ::NaturalParametersManifold{𝔽,TD,D,M,C,Nothing,FisherInformationMetric}, ::Type{T}
 ) where {𝔽,T,TD,D,M,C}
     return FirstOrderRetraction()
 end
@@ -141,3 +182,30 @@ function ManifoldsBase.retract!(
 )
     return ManifoldsBase.retract_fused!(M, q, p, X, one(eltype(X)), method)
 end
+
+struct NaturalBasis{𝔽,VST<:VectorSpaceType} <: AbstractBasis{𝔽,VST}
+    vector_space::VST
+end
+
+NaturalBasis(𝔽 = ℝ, vs::VectorSpaceType = TangentSpaceType()) = NaturalBasis{𝔽,typeof(vs)}(vs)
+NaturalBasis{𝔽}(vs::VectorSpaceType = TangentSpaceType()) where {𝔽} = NaturalBasis{𝔽,typeof(vs)}(vs)
+
+function ManifoldsBase.get_basis_default(
+    M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p
+) where {𝔽,T,D,MB,C,R}
+    return NaturalBasis{𝔽}()
+end
+
+function Manifolds.local_metric(
+    M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p, ::NaturalBasis
+) where {𝔽,T,D,MB,C,R}
+    ef = convert(ExponentialFamilyDistribution, M, p)
+    return ExponentialFamily.fisherinformation(ef)
+end
+
+function Manifolds.inverse_local_metric(
+    M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p, ::NaturalBasis
+) where {𝔽,T,D,MB,C,R}
+    ef = convert(ExponentialFamilyDistribution, M, p)
+    return cholinv(ExponentialFamily.fisherinformation(ef))
+end
diff --git a/src/natural_manifolds/normal.jl b/src/natural_manifolds/normal.jl
index e9927f9..e373e1b 100644
--- a/src/natural_manifolds/normal.jl
+++ b/src/natural_manifolds/normal.jl
@@ -68,7 +68,11 @@ function partition_point(
 end
 
 function ManifoldsBase.default_retraction_method(
-    ::NaturalParametersManifold{𝔽,MvNormalMeanCovariance,D,M,C,Nothing}, ::Type{T}
+    ::NaturalParametersManifold{𝔽,MvNormalMeanCovariance,D,M,C,Nothing,BaseMetric}, ::Type{T}
 ) where {𝔽,T,D,M,C}
     return ManifoldsBase.ExponentialRetraction()
 end
+
+function getdefaultmetric(::Type{MvNormalMeanCovariance})
+    return BaseMetric()
+end
diff --git a/test/natural_manifolds/natural_manifolds_setuptests.jl b/test/natural_manifolds/natural_manifolds_setuptests.jl
index 6ebb980..48ecd4f 100644
--- a/test/natural_manifolds/natural_manifolds_setuptests.jl
+++ b/test/natural_manifolds/natural_manifolds_setuptests.jl
@@ -1,9 +1,9 @@
 
-using StableRNGs, ExponentialFamily, ManifoldsBase, LinearAlgebra
+using StableRNGs, ExponentialFamily, Manifolds, ManifoldsBase, LinearAlgebra
 
 import ExponentialFamilyManifolds: get_natural_manifold, partition_point
 
-function test_natural_manifold(f; seed=42, ndistributions=100)
+function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
     rng = StableRNG(seed)
 
     foreach(1:ndistributions) do _
@@ -17,5 +17,11 @@ function test_natural_manifold(f; seed=42, ndistributions=100)
         η = partition_point(T, dims, getnaturalparameters(ef), getconditioner(ef))
 
         @test is_point(M, η, error=:error)
+
+        if test_metric
+            @test M.metric isa ExponentialFamilyManifolds.FisherInformationMetric
+            @test ManifoldsBase.get_basis_default(M, η) isa ExponentialFamilyManifolds.NaturalBasis
+            @test Manifolds.local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈ ExponentialFamily.fisherinformation(ef)
+        end
     end
 end
diff --git a/test/natural_manifolds/normal_tests.jl b/test/natural_manifolds/normal_tests.jl
index 561fc99..7cdba95 100644
--- a/test/natural_manifolds/normal_tests.jl
+++ b/test/natural_manifolds/normal_tests.jl
@@ -9,13 +9,21 @@ end
 @testitem "Check `MvNormal` natural manifold" begin
     include("natural_manifolds_setuptests.jl")
 
-    test_natural_manifold() do rng
+    test_natural_manifold(; test_metric=false) do rng
         k = rand(rng, 1:10)
         m = randn(k)
         L = LowerTriangular(randn(k, k))
         C = L * L' + k * I
         return MvNormalMeanCovariance(m, C)
     end
+
+    using ExponentialFamilyManifolds
+    using ExponentialFamily
+
+    @test ExponentialFamilyManifolds.getdefaultmetric(MvNormalMeanCovariance) isa
+        ExponentialFamilyManifolds.BaseMetric
+    M = ExponentialFamilyManifolds.get_natural_manifold(MvNormalMeanCovariance, (3,))
+    @test M.metric isa ExponentialFamilyManifolds.BaseMetric
 end
 
 @testitem "Check `MvNormalMeanScalePrecision` natural manifold" begin

From a1d7584a1c25534dc4e088182741d58990a2e8d8 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Mon, 17 Mar 2025 15:18:21 +0100
Subject: [PATCH 3/8] feat: draft generic SecondOrderRetraction

---
 Project.toml                              |  4 ++-
 src/natural_manifolds.jl                  | 40 +++++++++++++++++++++++
 test/natural_manifolds/bernoulli_tests.jl | 17 ++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 37e8b5a..5acabb5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -18,6 +18,7 @@ Static = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
 ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd"
 
 [compat]
+ADTypes = "1.14.0"
 BayesBase = "1.3"
 ExponentialFamily = "2.0.0"
 FastCholesky = "1.3.1"
@@ -31,6 +32,7 @@ Static = "0.8, 1"
 julia = "1.10"
 
 [extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -42,4 +44,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Distributions", "JET", "Test", "ReTestItems", "StableRNGs", "StaticArrays", "Manopt", "ForwardDiff"]
+test = ["ADTypes", "Aqua", "Distributions", "JET", "Test", "ReTestItems", "StableRNGs", "StaticArrays", "Manopt", "ForwardDiff"]
diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl
index 037bf0b..55f958e 100644
--- a/src/natural_manifolds.jl
+++ b/src/natural_manifolds.jl
@@ -59,6 +59,10 @@ function select_skip_methods(f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,F
         ManifoldsBase.retract_fused!,
         Manifolds.local_metric,
         Manifolds.local_metric_jacobian,
+        Manifolds.default_retraction_method,
+        Manifolds.get_basis_default,
+        Manifolds.christoffel_symbols_second,
+        Manifolds.christoffel_symbols_first
     )
         return ManifoldsBase.EmptyTrait()
     else
@@ -157,6 +161,21 @@ function ChartNOrderRetraction{O}() where {O}
 end
 
 const FirstOrderRetraction = ChartNOrderRetraction{1}
+const SecondOrderRetraction = ChartNOrderRetraction{2}
+
+"""
+    SecondOrderRetraction(; backend=nothing)
+
+Create a second-order retraction method that uses Christoffel symbols to compute
+a more accurate retraction. If a backend is provided, it will be used for any
+automatic differentiation needed to compute the Christoffel symbols.
+
+# Arguments
+- `backend`: Optional backend for automatic differentiation (e.g., `ADTypes.AutoForwardDiff()`)
+"""
+function SecondOrderRetraction(; backend=nothing)
+    return ChartNOrderRetraction{2,typeof(backend)}(backend)
+end
 
 function ManifoldsBase.default_retraction_method(
     ::NaturalParametersManifold{𝔽,TD,D,M,C,Nothing,FisherInformationMetric}, ::Type{T}
@@ -183,6 +202,27 @@ function ManifoldsBase.retract!(
     return ManifoldsBase.retract_fused!(M, q, p, X, one(eltype(X)), method)
 end
 
+function ManifoldsBase.retract_fused!(
+    M::NaturalParametersManifold{𝔽,T,D,BM,C,R,FisherInformationMetric}, 
+    q, p, X, t::Number, 
+    method::SecondOrderRetraction
+) where {𝔽,T,D,BM,C,R}
+    basis = ManifoldsBase.get_basis_default(M, p)
+    Γ = Manifolds.christoffel_symbols_second(M, p, basis; backend=method.extra)
+    
+    Δ = similar(p)
+    Manifolds.@einsum Δ[k] = -0.5 * Γ[k,i,j] * (t * X[i]) * (t * X[j])
+    
+    q .= p .+ t .* X .+ Δ
+    return q
+end
+
+function ManifoldsBase.retract!(
+    M::NaturalParametersManifold, q, p, X, method::SecondOrderRetraction
+)
+    return ManifoldsBase.retract_fused!(M, q, p, X, one(eltype(X)), method)
+end
+
 struct NaturalBasis{𝔽,VST<:VectorSpaceType} <: AbstractBasis{𝔽,VST}
     vector_space::VST
 end
diff --git a/test/natural_manifolds/bernoulli_tests.jl b/test/natural_manifolds/bernoulli_tests.jl
index 82c7292..82de212 100644
--- a/test/natural_manifolds/bernoulli_tests.jl
+++ b/test/natural_manifolds/bernoulli_tests.jl
@@ -5,3 +5,20 @@
         return Bernoulli(rand(rng))
     end
 end
+
+@testitem "Check SecondOrderRetraction" begin
+    include("natural_manifolds_setuptests.jl")
+
+    using ADTypes: AutoForwardDiff
+
+    rng = StableRNG(42)
+    M = ExponentialFamilyManifolds.get_natural_manifold(Bernoulli, ())
+    p = rand(rng, M)
+    X = rand(rng, M)
+    basis = ExponentialFamilyManifolds.NaturalBasis()
+
+    @show Manifolds.local_metric(M, p, basis)
+    @show Manifolds.local_metric_jacobian(M, p, basis, backend = AutoForwardDiff())
+    q = retract(M, p, X, ExponentialFamilyManifolds.SecondOrderRetraction(backend=AutoForwardDiff()))
+    @show q
+end
\ No newline at end of file

From 47c199700614d09fdf13221f5cc883242dc240e8 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Mon, 17 Mar 2025 15:18:42 +0100
Subject: [PATCH 4/8] style: :art:

---
 src/natural_manifolds.jl                      | 31 ++++++++++++-------
 src/natural_manifolds/normal.jl               |  3 +-
 test/natural_manifolds/bernoulli_tests.jl     | 11 +++++--
 .../natural_manifolds_setuptests.jl           |  6 ++--
 4 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl
index 55f958e..c7c7888 100644
--- a/src/natural_manifolds.jl
+++ b/src/natural_manifolds.jl
@@ -47,11 +47,15 @@ getretraction(M::NaturalParametersManifold) = M.retraction
 
 # The `NaturalParametersManifold` simply adds extra properties to the `base` and 
 # acts as a "decorator"
-function select_skip_methods(::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,BaseMetric}) where {F,𝔽,T,D,MB,C,R}
+function select_skip_methods(
+    ::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,BaseMetric}
+) where {F,𝔽,T,D,MB,C,R}
     return ManifoldsBase.IsExplicitDecorator()
 end
 
-function select_skip_methods(f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}) where {F,𝔽,T,D,MB,C,R}
+function select_skip_methods(
+    f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}
+) where {F,𝔽,T,D,MB,C,R}
     if f in (
         ManifoldsBase.retract,
         ManifoldsBase.retract!,
@@ -62,7 +66,7 @@ function select_skip_methods(f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,F
         Manifolds.default_retraction_method,
         Manifolds.get_basis_default,
         Manifolds.christoffel_symbols_second,
-        Manifolds.christoffel_symbols_first
+        Manifolds.christoffel_symbols_first,
     )
         return ManifoldsBase.EmptyTrait()
     else
@@ -203,16 +207,19 @@ function ManifoldsBase.retract!(
 end
 
 function ManifoldsBase.retract_fused!(
-    M::NaturalParametersManifold{𝔽,T,D,BM,C,R,FisherInformationMetric}, 
-    q, p, X, t::Number, 
-    method::SecondOrderRetraction
+    M::NaturalParametersManifold{𝔽,T,D,BM,C,R,FisherInformationMetric},
+    q,
+    p,
+    X,
+    t::Number,
+    method::SecondOrderRetraction,
 ) where {𝔽,T,D,BM,C,R}
     basis = ManifoldsBase.get_basis_default(M, p)
     Γ = Manifolds.christoffel_symbols_second(M, p, basis; backend=method.extra)
-    
+
     Δ = similar(p)
-    Manifolds.@einsum Δ[k] = -0.5 * Γ[k,i,j] * (t * X[i]) * (t * X[j])
-    
+    Manifolds.@einsum Δ[k] = -0.5 * Γ[k, i, j] * (t * X[i]) * (t * X[j])
+
     q .= p .+ t .* X .+ Δ
     return q
 end
@@ -227,8 +234,10 @@ struct NaturalBasis{𝔽,VST<:VectorSpaceType} <: AbstractBasis{𝔽,VST}
     vector_space::VST
 end
 
-NaturalBasis(𝔽 = ℝ, vs::VectorSpaceType = TangentSpaceType()) = NaturalBasis{𝔽,typeof(vs)}(vs)
-NaturalBasis{𝔽}(vs::VectorSpaceType = TangentSpaceType()) where {𝔽} = NaturalBasis{𝔽,typeof(vs)}(vs)
+NaturalBasis(𝔽=ℝ, vs::VectorSpaceType=TangentSpaceType()) = NaturalBasis{𝔽,typeof(vs)}(vs)
+function NaturalBasis{𝔽}(vs::VectorSpaceType=TangentSpaceType()) where {𝔽}
+    return NaturalBasis{𝔽,typeof(vs)}(vs)
+end
 
 function ManifoldsBase.get_basis_default(
     M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p
diff --git a/src/natural_manifolds/normal.jl b/src/natural_manifolds/normal.jl
index e373e1b..943f554 100644
--- a/src/natural_manifolds/normal.jl
+++ b/src/natural_manifolds/normal.jl
@@ -68,7 +68,8 @@ function partition_point(
 end
 
 function ManifoldsBase.default_retraction_method(
-    ::NaturalParametersManifold{𝔽,MvNormalMeanCovariance,D,M,C,Nothing,BaseMetric}, ::Type{T}
+    ::NaturalParametersManifold{𝔽,MvNormalMeanCovariance,D,M,C,Nothing,BaseMetric},
+    ::Type{T},
 ) where {𝔽,T,D,M,C}
     return ManifoldsBase.ExponentialRetraction()
 end
diff --git a/test/natural_manifolds/bernoulli_tests.jl b/test/natural_manifolds/bernoulli_tests.jl
index 82de212..f382037 100644
--- a/test/natural_manifolds/bernoulli_tests.jl
+++ b/test/natural_manifolds/bernoulli_tests.jl
@@ -18,7 +18,12 @@ end
     basis = ExponentialFamilyManifolds.NaturalBasis()
 
     @show Manifolds.local_metric(M, p, basis)
-    @show Manifolds.local_metric_jacobian(M, p, basis, backend = AutoForwardDiff())
-    q = retract(M, p, X, ExponentialFamilyManifolds.SecondOrderRetraction(backend=AutoForwardDiff()))
+    @show Manifolds.local_metric_jacobian(M, p, basis, backend=AutoForwardDiff())
+    q = retract(
+        M,
+        p,
+        X,
+        ExponentialFamilyManifolds.SecondOrderRetraction(; backend=AutoForwardDiff()),
+    )
     @show q
-end
\ No newline at end of file
+end
diff --git a/test/natural_manifolds/natural_manifolds_setuptests.jl b/test/natural_manifolds/natural_manifolds_setuptests.jl
index 48ecd4f..942f93d 100644
--- a/test/natural_manifolds/natural_manifolds_setuptests.jl
+++ b/test/natural_manifolds/natural_manifolds_setuptests.jl
@@ -20,8 +20,10 @@ function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
 
         if test_metric
             @test M.metric isa ExponentialFamilyManifolds.FisherInformationMetric
-            @test ManifoldsBase.get_basis_default(M, η) isa ExponentialFamilyManifolds.NaturalBasis
-            @test Manifolds.local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈ ExponentialFamily.fisherinformation(ef)
+            @test ManifoldsBase.get_basis_default(M, η) isa
+                ExponentialFamilyManifolds.NaturalBasis
+            @test Manifolds.local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈
+                ExponentialFamily.fisherinformation(ef)
         end
     end
 end

From 4d6fad449822131531e3db7820ab9a47ab1ca773 Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 18 Mar 2025 10:56:06 +0100
Subject: [PATCH 5/8] refactor: use metric to select retraction method

---
 src/natural_manifolds.jl                      | 112 +++++++++---------
 src/natural_manifolds/normal.jl               |   7 --
 .../natural_manifolds_setuptests.jl           |   4 +
 test/natural_manifolds_tests.jl               |  14 +++
 4 files changed, 75 insertions(+), 62 deletions(-)

diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl
index c7c7888..192976e 100644
--- a/src/natural_manifolds.jl
+++ b/src/natural_manifolds.jl
@@ -3,12 +3,45 @@ using FastCholesky
 
 import ExponentialFamily: exponential_family_typetag
 
+
+struct ChartNOrderRetraction{Order,E} <: AbstractRetractionMethod
+    extra::E
+end
+
+function ChartNOrderRetraction{O}() where {O}
+    return ChartNOrderRetraction{O,Nothing}(nothing)
+end
+
+const FirstOrderRetraction = ChartNOrderRetraction{1}
+const SecondOrderRetraction = ChartNOrderRetraction{2}
+
+"""
+    SecondOrderRetraction(; backend=nothing)
+
+Create a second-order retraction method that uses Christoffel symbols to compute
+a more accurate retraction. If a backend is provided, it will be used for any
+automatic differentiation needed to compute the Christoffel symbols.
+
+# Arguments
+- `backend`: Optional backend for automatic differentiation (e.g., `ADTypes.AutoForwardDiff()`)
+"""
+function SecondOrderRetraction(; backend=nothing)
+    return ChartNOrderRetraction{2,typeof(backend)}(backend)
+end
+
 """
     FisherInformationMetric <: RiemannianMetric
 
 Specifier that we need to use the Fisher information metric.
 """
-struct FisherInformationMetric <: RiemannianMetric end
+struct FisherInformationMetric{R} <: RiemannianMetric 
+    default_retraction::R
+end
+
+function FisherInformationMetric()
+    retraction = FirstOrderRetraction()
+    return FisherInformationMetric{typeof(retraction)}(retraction)
+end
 
 """
     BaseMetric <: RiemannianMetric
@@ -32,30 +65,29 @@ end
 The manifold for the natural parameters of the distribution of type `T` with dimensions `dims`.
 An internal structure, use `get_natural_manifold` to create an instance of a manifold for the natural parameters of distribution of type `T`.
 """
-struct NaturalParametersManifold{𝔽,T,D,M,C,R,MT} <: AbstractDecoratorManifold{𝔽}
+struct NaturalParametersManifold{𝔽,T,D,M,C,MT} <: AbstractDecoratorManifold{𝔽}
     dims::D
     base::M
     conditioner::C
-    retraction::R
     metric::MT
 end
 
 getdims(M::NaturalParametersManifold) = M.dims
 getbase(M::NaturalParametersManifold) = M.base
 getconditioner(M::NaturalParametersManifold) = M.conditioner
-getretraction(M::NaturalParametersManifold) = M.retraction
+getmetric(M::NaturalParametersManifold) = M.metric
 
 # The `NaturalParametersManifold` simply adds extra properties to the `base` and 
 # acts as a "decorator"
 function select_skip_methods(
-    ::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,BaseMetric}
-) where {F,𝔽,T,D,MB,C,R}
+    ::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,BaseMetric}
+) where {F,𝔽,T,D,MB,C}
     return ManifoldsBase.IsExplicitDecorator()
 end
 
 function select_skip_methods(
-    f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}
-) where {F,𝔽,T,D,MB,C,R}
+    f::F, ::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}
+) where {F,𝔽,T,D,MB,C}
     if f in (
         ManifoldsBase.retract,
         ManifoldsBase.retract!,
@@ -63,6 +95,7 @@ function select_skip_methods(
         ManifoldsBase.retract_fused!,
         Manifolds.local_metric,
         Manifolds.local_metric_jacobian,
+        Manifolds.inverse_local_metric,
         Manifolds.default_retraction_method,
         Manifolds.get_basis_default,
         Manifolds.christoffel_symbols_second,
@@ -93,11 +126,10 @@ function NaturalParametersManifold(
     dims::D,
     base::M,
     conditioner::C=nothing,
-    retraction::R=nothing,
     metric::MT=getdefaultmetric(T),
-) where {T,𝔽,D,M<:AbstractManifold{𝔽},C,R,MT}
-    return NaturalParametersManifold{𝔽,T,D,M,C,R,MT}(
-        dims, base, conditioner, retraction, metric
+) where {T,𝔽,D,M<:AbstractManifold{𝔽},C,MT}
+    return NaturalParametersManifold{𝔽,T,D,M,C,MT}(
+        dims, base, conditioner, metric
     )
 end
 
@@ -118,10 +150,10 @@ true
 ```
 """
 function get_natural_manifold(
-    ::Type{T}, dims, conditioner=nothing, retraction=nothing
+    ::Type{T}, dims, conditioner=nothing, metric=getdefaultmetric(T)
 ) where {T}
     return NaturalParametersManifold(
-        T, dims, get_natural_manifold_base(T, dims, conditioner), conditioner, retraction
+        T, dims, get_natural_manifold_base(T, dims, conditioner), conditioner, metric
     )
 end
 
@@ -156,41 +188,11 @@ function Base.convert(
     )
 end
 
-struct ChartNOrderRetraction{Order,E} <: AbstractRetractionMethod
-    extra::E
-end
-
-function ChartNOrderRetraction{O}() where {O}
-    return ChartNOrderRetraction{O,Nothing}(nothing)
-end
-
-const FirstOrderRetraction = ChartNOrderRetraction{1}
-const SecondOrderRetraction = ChartNOrderRetraction{2}
-
-"""
-    SecondOrderRetraction(; backend=nothing)
-
-Create a second-order retraction method that uses Christoffel symbols to compute
-a more accurate retraction. If a backend is provided, it will be used for any
-automatic differentiation needed to compute the Christoffel symbols.
-
-# Arguments
-- `backend`: Optional backend for automatic differentiation (e.g., `ADTypes.AutoForwardDiff()`)
-"""
-function SecondOrderRetraction(; backend=nothing)
-    return ChartNOrderRetraction{2,typeof(backend)}(backend)
-end
-
-function ManifoldsBase.default_retraction_method(
-    ::NaturalParametersManifold{𝔽,TD,D,M,C,Nothing,FisherInformationMetric}, ::Type{T}
-) where {𝔽,T,TD,D,M,C}
-    return FirstOrderRetraction()
-end
 
 function ManifoldsBase.default_retraction_method(
-    M::NaturalParametersManifold{𝔽,TD,D,BM,C,R}, ::Type{T}
-) where {𝔽,T,TD,D,BM,C,R}
-    return getretraction(M)
+    M::NaturalParametersManifold{𝔽,TD,D,BM,C,<:FisherInformationMetric}, ::Type{T}
+) where {𝔽,T,TD,D,BM,C}
+    return getmetric(M).default_retraction
 end
 
 function ManifoldsBase.retract_fused!(
@@ -207,13 +209,13 @@ function ManifoldsBase.retract!(
 end
 
 function ManifoldsBase.retract_fused!(
-    M::NaturalParametersManifold{𝔽,T,D,BM,C,R,FisherInformationMetric},
+    M::NaturalParametersManifold{𝔽,T,D,BM,C,<:FisherInformationMetric},
     q,
     p,
     X,
     t::Number,
     method::SecondOrderRetraction,
-) where {𝔽,T,D,BM,C,R}
+) where {𝔽,T,D,BM,C}
     basis = ManifoldsBase.get_basis_default(M, p)
     Γ = Manifolds.christoffel_symbols_second(M, p, basis; backend=method.extra)
 
@@ -240,21 +242,21 @@ function NaturalBasis{𝔽}(vs::VectorSpaceType=TangentSpaceType()) where {𝔽}
 end
 
 function ManifoldsBase.get_basis_default(
-    M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p
-) where {𝔽,T,D,MB,C,R}
+    ::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}, p
+) where {𝔽,T,D,MB,C}
     return NaturalBasis{𝔽}()
 end
 
 function Manifolds.local_metric(
-    M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p, ::NaturalBasis
-) where {𝔽,T,D,MB,C,R}
+    M::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}, p, ::NaturalBasis
+) where {𝔽,T,D,MB,C}
     ef = convert(ExponentialFamilyDistribution, M, p)
     return ExponentialFamily.fisherinformation(ef)
 end
 
 function Manifolds.inverse_local_metric(
-    M::NaturalParametersManifold{𝔽,T,D,MB,C,R,FisherInformationMetric}, p, ::NaturalBasis
-) where {𝔽,T,D,MB,C,R}
+    M::NaturalParametersManifold{𝔽,T,D,MB,C,<:FisherInformationMetric}, p, ::NaturalBasis
+) where {𝔽,T,D,MB,C}
     ef = convert(ExponentialFamilyDistribution, M, p)
     return cholinv(ExponentialFamily.fisherinformation(ef))
 end
diff --git a/src/natural_manifolds/normal.jl b/src/natural_manifolds/normal.jl
index 943f554..aec181d 100644
--- a/src/natural_manifolds/normal.jl
+++ b/src/natural_manifolds/normal.jl
@@ -67,13 +67,6 @@ function partition_point(
     return ArrayPartition(view(p, 1:k), view(p, (k + 1):(k + 1)))
 end
 
-function ManifoldsBase.default_retraction_method(
-    ::NaturalParametersManifold{𝔽,MvNormalMeanCovariance,D,M,C,Nothing,BaseMetric},
-    ::Type{T},
-) where {𝔽,T,D,M,C}
-    return ManifoldsBase.ExponentialRetraction()
-end
-
 function getdefaultmetric(::Type{MvNormalMeanCovariance})
     return BaseMetric()
 end
diff --git a/test/natural_manifolds/natural_manifolds_setuptests.jl b/test/natural_manifolds/natural_manifolds_setuptests.jl
index 942f93d..2e1e5c3 100644
--- a/test/natural_manifolds/natural_manifolds_setuptests.jl
+++ b/test/natural_manifolds/natural_manifolds_setuptests.jl
@@ -3,6 +3,8 @@ using StableRNGs, ExponentialFamily, Manifolds, ManifoldsBase, LinearAlgebra
 
 import ExponentialFamilyManifolds: get_natural_manifold, partition_point
 
+using FastCholesky
+
 function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
     rng = StableRNG(seed)
 
@@ -24,6 +26,8 @@ function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
                 ExponentialFamilyManifolds.NaturalBasis
             @test Manifolds.local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈
                 ExponentialFamily.fisherinformation(ef)
+            @test Manifolds.inverse_local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈
+                cholinv(ExponentialFamily.fisherinformation(ef))
         end
     end
 end
diff --git a/test/natural_manifolds_tests.jl b/test/natural_manifolds_tests.jl
index 2342938..b230f71 100644
--- a/test/natural_manifolds_tests.jl
+++ b/test/natural_manifolds_tests.jl
@@ -76,3 +76,17 @@
         end
     end
 end
+
+@testitem "Natural manifold properties: BaseMetric" begin
+    using Distributions, ExponentialFamily, Manifolds, ManifoldsBase, StableRNGs, JET, LinearAlgebra
+
+    import ExponentialFamilyManifolds: BaseMetric
+
+    M = ExponentialFamilyManifolds.get_natural_manifold(Beta, (), nothing, BaseMetric())
+    @test M.metric isa BaseMetric
+    @test ExponentialFamilyManifolds.select_skip_methods(ManifoldsBase.retract, M) == ManifoldsBase.IsExplicitDecorator()
+    p = rand(StableRNG(42), M)
+    q = copy(p)
+    @test  ManifoldsBase.retract!(M, q, p, p) ≈ ManifoldsBase.retract!(M.base, q, p, p)
+end
+

From 9981d6e661d259f25b6bc98d64d47d1c8397ad9a Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 18 Mar 2025 12:21:48 +0100
Subject: [PATCH 6/8] test: first order retraction

---
 src/natural_manifolds.jl                      |  2 +-
 src/natural_manifolds/beta.jl                 |  2 ++
 src/natural_manifolds/categorical.jl          |  4 +++
 src/natural_manifolds/dirichlet.jl            |  5 ++++
 src/natural_manifolds/gamma.jl                |  2 ++
 src/natural_manifolds/inverse_gamma.jl        |  2 ++
 src/natural_manifolds/lognormal.jl            |  2 ++
 src/natural_manifolds/normal.jl               |  6 ++++
 src/natural_manifolds/wishart.jl              |  4 +++
 test/natural_manifolds/beta_tests.jl          |  1 +
 test/natural_manifolds/categorical_tests.jl   |  2 +-
 .../natural_manifolds_setuptests.jl           | 30 +++++++++++++++++--
 12 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/src/natural_manifolds.jl b/src/natural_manifolds.jl
index 192976e..29f5667 100644
--- a/src/natural_manifolds.jl
+++ b/src/natural_manifolds.jl
@@ -100,6 +100,7 @@ function select_skip_methods(
         Manifolds.get_basis_default,
         Manifolds.christoffel_symbols_second,
         Manifolds.christoffel_symbols_first,
+        Manifolds.representation_size
     )
         return ManifoldsBase.EmptyTrait()
     else
@@ -221,7 +222,6 @@ function ManifoldsBase.retract_fused!(
 
     Δ = similar(p)
     Manifolds.@einsum Δ[k] = -0.5 * Γ[k, i, j] * (t * X[i]) * (t * X[j])
-
     q .= p .+ t .* X .+ Δ
     return q
 end
diff --git a/src/natural_manifolds/beta.jl b/src/natural_manifolds/beta.jl
index a9849f6..df949f8 100644
--- a/src/natural_manifolds/beta.jl
+++ b/src/natural_manifolds/beta.jl
@@ -17,3 +17,5 @@ Converts the `point` to a compatible representation for the natural manifold of
 function partition_point(::Type{Beta}, ::Tuple{}, p, conditioner=nothing)
     return ArrayPartition(view(p, 1:1), view(p, 2:2))
 end
+
+Manifolds.representation_size(::NaturalParametersManifold{𝔽, Beta}) where {𝔽} = (2,)
\ No newline at end of file
diff --git a/src/natural_manifolds/categorical.jl b/src/natural_manifolds/categorical.jl
index f35b8e2..5de7736 100644
--- a/src/natural_manifolds/categorical.jl
+++ b/src/natural_manifolds/categorical.jl
@@ -16,3 +16,7 @@ Converts the `point` to a compatible representation for the natural manifold of
 function partition_point(::Type{Categorical}, ::Tuple{}, p, conditioner=nothing)
     return ArrayPartition(view(p, 1:(conditioner - 1)), view(p, conditioner:conditioner))
 end
+
+function Manifolds.representation_size(M::NaturalParametersManifold{𝔽, Categorical}) where {𝔽}
+    return (getconditioner(M),)
+end
diff --git a/src/natural_manifolds/dirichlet.jl b/src/natural_manifolds/dirichlet.jl
index dd0e230..58f2f96 100644
--- a/src/natural_manifolds/dirichlet.jl
+++ b/src/natural_manifolds/dirichlet.jl
@@ -19,3 +19,8 @@ function partition_point(::Type{Dirichlet}, dims::Tuple{Int}, p, conditioner=not
     # See comment in `get_natural_manifold_base` for `Dirichlet`
     return ArrayPartition(p')
 end
+
+function Manifolds.representation_size(M::NaturalParametersManifold{𝔽, Dirichlet}) where {𝔽}
+    dims = getdims(M)
+    return (first(dims),)
+end
\ No newline at end of file
diff --git a/src/natural_manifolds/gamma.jl b/src/natural_manifolds/gamma.jl
index d910c70..412f3bc 100644
--- a/src/natural_manifolds/gamma.jl
+++ b/src/natural_manifolds/gamma.jl
@@ -18,3 +18,5 @@ Converts the `point` to a compatible representation for the natural manifold of
 function partition_point(::Type{Gamma}, ::Tuple{}, p, conditioner=nothing)
     return ArrayPartition(view(p, 1:1), view(p, 2:2))
 end
+
+Manifolds.representation_size(::NaturalParametersManifold{𝔽, Gamma}) where {𝔽} = (2,)
\ No newline at end of file
diff --git a/src/natural_manifolds/inverse_gamma.jl b/src/natural_manifolds/inverse_gamma.jl
index d046e2b..57734dd 100644
--- a/src/natural_manifolds/inverse_gamma.jl
+++ b/src/natural_manifolds/inverse_gamma.jl
@@ -21,3 +21,5 @@ function partition_point(
 )
     return ArrayPartition(view(p, 1:1), view(p, 2:2))
 end
+
+Manifolds.representation_size(::NaturalParametersManifold{𝔽, ExponentialFamily.GammaInverse}) where {𝔽} = (2,)
\ No newline at end of file
diff --git a/src/natural_manifolds/lognormal.jl b/src/natural_manifolds/lognormal.jl
index 90c7104..67e6532 100644
--- a/src/natural_manifolds/lognormal.jl
+++ b/src/natural_manifolds/lognormal.jl
@@ -15,3 +15,5 @@ Converts the `point` to a compatible representation for the natural manifold of
 function partition_point(::Type{LogNormal}, ::Tuple{}, p, conditioner=nothing)
     return ArrayPartition(view(p, 1:1), view(p, 2:2))
 end
+
+Manifolds.representation_size(::NaturalParametersManifold{𝔽, LogNormal}) where {𝔽} = (2,)
\ No newline at end of file
diff --git a/src/natural_manifolds/normal.jl b/src/natural_manifolds/normal.jl
index aec181d..d0a0321 100644
--- a/src/natural_manifolds/normal.jl
+++ b/src/natural_manifolds/normal.jl
@@ -70,3 +70,9 @@ end
 function getdefaultmetric(::Type{MvNormalMeanCovariance})
     return BaseMetric()
 end
+
+Manifolds.representation_size(::NaturalParametersManifold{𝔽, NormalMeanVariance}) where {𝔽} = (2,)
+
+function Manifolds.representation_size(M::NaturalParametersManifold{𝔽, MvNormalMeanScalePrecision}) where {𝔽}
+    return (M.dims[1] +1, )
+end
\ No newline at end of file
diff --git a/src/natural_manifolds/wishart.jl b/src/natural_manifolds/wishart.jl
index 0749200..5269d17 100644
--- a/src/natural_manifolds/wishart.jl
+++ b/src/natural_manifolds/wishart.jl
@@ -21,3 +21,7 @@ function partition_point(
     k = first(dims)
     return ArrayPartition(view(p, 1:1), reshape(view(p, 2:(1 + k^2)), (k, k)))
 end
+
+function getdefaultmetric(::Type{ExponentialFamily.WishartFast})
+    return BaseMetric()
+end
\ No newline at end of file
diff --git a/test/natural_manifolds/beta_tests.jl b/test/natural_manifolds/beta_tests.jl
index 5df840e..7d2b68a 100644
--- a/test/natural_manifolds/beta_tests.jl
+++ b/test/natural_manifolds/beta_tests.jl
@@ -4,4 +4,5 @@
     test_natural_manifold() do rng
         return Beta(10rand(rng), 10rand(rng))
     end
+
 end
diff --git a/test/natural_manifolds/categorical_tests.jl b/test/natural_manifolds/categorical_tests.jl
index dce5b2e..053bf61 100644
--- a/test/natural_manifolds/categorical_tests.jl
+++ b/test/natural_manifolds/categorical_tests.jl
@@ -1,7 +1,7 @@
 @testitem "Check `Categorical` natural manifold" begin
     include("natural_manifolds_setuptests.jl")
 
-    test_natural_manifold() do rng
+    test_natural_manifold(test_injectivity_radius=false) do rng
         p = rand(rng, 10)
         normalize!(p, 1)
         return Categorical(p)
diff --git a/test/natural_manifolds/natural_manifolds_setuptests.jl b/test/natural_manifolds/natural_manifolds_setuptests.jl
index 2e1e5c3..003d810 100644
--- a/test/natural_manifolds/natural_manifolds_setuptests.jl
+++ b/test/natural_manifolds/natural_manifolds_setuptests.jl
@@ -5,7 +5,7 @@ import ExponentialFamilyManifolds: get_natural_manifold, partition_point
 
 using FastCholesky
 
-function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
+function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true, test_injectivity_radius=true)
     rng = StableRNG(seed)
 
     foreach(1:ndistributions) do _
@@ -20,8 +20,7 @@ function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
 
         @test is_point(M, η, error=:error)
 
-        if test_metric
-            @test M.metric isa ExponentialFamilyManifolds.FisherInformationMetric
+        if test_metric && M.metric isa ExponentialFamilyManifolds.FisherInformationMetric
             @test ManifoldsBase.get_basis_default(M, η) isa
                 ExponentialFamilyManifolds.NaturalBasis
             @test Manifolds.local_metric(M, η, ManifoldsBase.get_basis_default(M, η)) ≈
@@ -30,4 +29,29 @@ function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true)
                 cholinv(ExponentialFamily.fisherinformation(ef))
         end
     end
+
+    distribution = f(rng)
+    sample = rand(rng, distribution)
+    dims = size(sample)
+    ef = convert(ExponentialFamilyDistribution, distribution)
+
+    T = ExponentialFamily.exponential_family_typetag(ef)
+    M = get_natural_manifold(T, dims, getconditioner(ef))
+    if test_metric && M.metric isa ExponentialFamilyManifolds.FisherInformationMetric
+        @testset "Testing $(typeof(M)) with retractions" begin
+            pts = [rand(rng, M) for _ in 1:5]
+            Manifolds.test_manifold(
+                M, 
+                pts;
+                test_exp_log=true,
+                default_inverse_retraction_method=nothing,
+                test_default_vector_transport=false,
+                retraction_methods=[
+                    ExponentialFamilyManifolds.FirstOrderRetraction(),
+                ],
+                inverse_retraction_methods=[],
+                test_injectivity_radius=test_injectivity_radius,
+            )
+        end
+    end
 end

From 621880e3e62d9a5d29cc0fbd4d24992aab14814d Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Tue, 18 Mar 2025 12:33:58 +0100
Subject: [PATCH 7/8] test: fisher of MvNormalMeanScalePrecision is not fd
 friendly

---
 .../natural_manifolds_setuptests.jl           | 21 ++++++++++++++-----
 test/natural_manifolds/normal_tests.jl        |  2 +-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/test/natural_manifolds/natural_manifolds_setuptests.jl b/test/natural_manifolds/natural_manifolds_setuptests.jl
index 003d810..de00742 100644
--- a/test/natural_manifolds/natural_manifolds_setuptests.jl
+++ b/test/natural_manifolds/natural_manifolds_setuptests.jl
@@ -1,11 +1,11 @@
 
 using StableRNGs, ExponentialFamily, Manifolds, ManifoldsBase, LinearAlgebra
-
+using ADTypes: AutoForwardDiff
 import ExponentialFamilyManifolds: get_natural_manifold, partition_point
 
 using FastCholesky
 
-function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true, test_injectivity_radius=true)
+function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true, test_injectivity_radius=true, fisher_fd_friendly=true)
     rng = StableRNG(seed)
 
     foreach(1:ndistributions) do _
@@ -40,18 +40,29 @@ function test_natural_manifold(f; seed=42, ndistributions=100, test_metric=true,
     if test_metric && M.metric isa ExponentialFamilyManifolds.FisherInformationMetric
         @testset "Testing $(typeof(M)) with retractions" begin
             pts = [rand(rng, M) for _ in 1:5]
+
+            if fisher_fd_friendly
+                retraction_methods = [
+                    ExponentialFamilyManifolds.FirstOrderRetraction(),
+                    ExponentialFamilyManifolds.SecondOrderRetraction(backend=AutoForwardDiff()),
+                ]
+            else
+                retraction_methods = [
+                    ExponentialFamilyManifolds.FirstOrderRetraction(),
+                ]
+            end
+
             Manifolds.test_manifold(
                 M, 
                 pts;
                 test_exp_log=true,
                 default_inverse_retraction_method=nothing,
                 test_default_vector_transport=false,
-                retraction_methods=[
-                    ExponentialFamilyManifolds.FirstOrderRetraction(),
-                ],
+                retraction_methods=retraction_methods,
                 inverse_retraction_methods=[],
                 test_injectivity_radius=test_injectivity_radius,
             )
         end
     end
 end
+  
\ No newline at end of file
diff --git a/test/natural_manifolds/normal_tests.jl b/test/natural_manifolds/normal_tests.jl
index 7cdba95..caa5fa7 100644
--- a/test/natural_manifolds/normal_tests.jl
+++ b/test/natural_manifolds/normal_tests.jl
@@ -29,7 +29,7 @@ end
 @testitem "Check `MvNormalMeanScalePrecision` natural manifold" begin
     include("natural_manifolds_setuptests.jl")
 
-    test_natural_manifold() do rng
+    test_natural_manifold(fisher_fd_friendly=false) do rng
         k = rand(rng, 1:10)
         m = randn(rng, k)
         γ = rand(rng)^2 + 1

From cb1b998e4c8e910e52d594a8c79275102cd66fff Mon Sep 17 00:00:00 2001
From: Mykola Lukashchuk <nikola.lukashuk@gmail.com>
Date: Wed, 19 Mar 2025 12:53:10 +0100
Subject: [PATCH 8/8] fix: update BayesBase for InvArrowHead

---
 Project.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Project.toml b/Project.toml
index 5acabb5..efef5b7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -19,7 +19,7 @@ ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd"
 
 [compat]
 ADTypes = "1.14.0"
-BayesBase = "1.3"
+BayesBase = "1.5.4"
 ExponentialFamily = "2.0.0"
 FastCholesky = "1.3.1"
 LinearAlgebra = "1.10"