Add manifold for Categorical distribution

src/ExponentialFamilyManifolds.jl

@@ -11,6 +11,7 @@ include("natural_manifolds/bernoulli.jl")
 include("natural_manifolds/beta.jl")
 include("natural_manifolds/binomial.jl")
 include("natural_manifolds/chisq.jl")
+include("natural_manifolds/categorical.jl")
 include("natural_manifolds/dirichlet.jl")
 include("natural_manifolds/exponential.jl")
 include("natural_manifolds/gamma.jl")

src/natural_manifolds/categorical.jl

@@ -0,0 +1,19 @@
+
+"""
+    get_natural_manifold_base(::Type{Categorical}, dims::Tuple{Int}, conditioner=nothing)
+
+Get the natural manifold base for the `Categorical` distribution.
+"""
+function get_natural_manifold_base(::Type{Categorical}, ::Tuple{}, conditioner=nothing)
+    return Euclidean(conditioner)
+end
+
+"""
+    partition_point(::Type{Categorical}, dims::Tuple{Int}, p, conditioner=nothing)
+
+Converts the `point` to a compatible representation for the natural manifold of type `Categorical`.
+"""
+function partition_point(::Type{Categorical}, ::Tuple{}, p, conditioner=nothing)
+    # See comment in `get_natural_manifold_base` for `Categorical`
+    return ArrayPartition(p)
+end

test/natural_manifolds/categorical_tests.jl

@@ -0,0 +1,9 @@
+@testitem "Check `Categorical` natural manifold" begin
+    include("natural_manifolds_setuptests.jl")
+
+    test_natural_manifold() do rng
+        p = rand(rng, 10)
+        normalize!(p, 1)
+        return Categorical(p)
+    end
+end