Dirichlet tensor

Project.toml

@@ -5,6 +5,7 @@ version = "1.6.0"
 
 [deps]
 BayesBase = "b4ee3484-f114-42fe-b91c-797d54a0c67e"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 FastCholesky = "2d5283b6-8564-42b6-bb00-83ed8e915756"
@@ -29,6 +30,7 @@ TinyHugeNumbers = "783c9a47-75a3-44ac-a16b-f1ab7b3acf04"
 [compat]
 Aqua = "0.8.7"
 BayesBase = "1.5.0"
+BlockArrays = "1.3.0"
 Distributions = "0.25"
 DomainSets = "0.5.2, 0.6, 0.7"
 FastCholesky = "1.0"

benchmark/benchmarks.jl

@@ -4,4 +4,5 @@ using ExponentialFamily.BayesBase
 
 const SUITE = BenchmarkGroup()
 
-include("benchmarks/bernoulli.jl")
+include("benchmarks/bernoulli.jl")
+include("benchmarks/tensordirichlet.jl")

benchmark/benchmarks/tensordirichlet.jl

@@ -0,0 +1,42 @@
+
+SUITE["tensordirichlet"] = BenchmarkGroup(
+    ["tensordirichlet", "distribution"],
+    "prod" => BenchmarkGroup(["prod", "multiplication"]),
+    "convert" => BenchmarkGroup(["convert"])
+)
+
+# `prod` BenchmarkGroup ========================
+for rank in (3, 4, 5, 6)
+    for d in (5, 10, 20)
+        left = TensorDirichlet(rand([d for _ in 1:rank]...) .+ 1)
+        right = TensorDirichlet(rand([d for _ in 1:rank]...) .+ 1)
+        SUITE["tensordirichlet"]["prod"]["Closed(rank=$rank, d=$d)"] = @benchmarkable prod(ClosedProd(), $left, $right)
+    end
+end
+
+# ==============================================
+
+# `convert` BenchmarkGroup =====================
+SUITE["tensordirichlet"]["convert"]["Convert from D to EF"] = @benchmarkable convert(ExponentialFamilyDistribution, dist) setup = begin
+    dist = TensorDirichlet(rand(5, 5, 5))
+end
+
+SUITE["tensordirichlet"]["convert"]["Convert from EF to D"] = @benchmarkable convert(Distribution, efdist) setup = begin
+    efdist = convert(ExponentialFamilyDistribution, TensorDirichlet(rand(5, 5, 5)))
+end
+# ==============================================
+
+for rank in (3, 4, 5, 6)
+    for d in (5, 10, 20)
+        distribution = TensorDirichlet(rand([d for _ in 1:rank]...))
+        sample = rand(distribution)
+        SUITE["tensordirichlet"]["mean"]["rank=$rank, d=$d"] = @benchmarkable mean($distribution)
+        SUITE["tensordirichlet"]["rand"]["rank=$rank, d=$d"] = @benchmarkable rand($distribution)
+        SUITE["tensordirichlet"]["logpdf"]["rank=$rank, d=$d"] = @benchmarkable logpdf($distribution, $sample)
+        SUITE["tensordirichlet"]["var"]["rank=$rank, d=$d"] = @benchmarkable var($distribution)
+        SUITE["tensordirichlet"]["cov"]["rank=$rank, d=$d"] = @benchmarkable cov($distribution)
+        SUITE["tensordirichlet"]["entropy"]["rank=$rank, d=$d"] = @benchmarkable entropy($distribution)
+    end
+end
+
+# ==============================================

docs/src/library.md

@@ -21,4 +21,5 @@ ExponentialFamily.WishartFast
 ExponentialFamily.InverseWishartFast
 ExponentialFamily.NormalGamma
 ExponentialFamily.MvNormalWishart
+ExponentialFamily.TensorDirichlet
 ```

scripts/benchmark.jl

@@ -10,6 +10,6 @@ if isempty(ARGS)
     export_markdown("./benchmark_logs/last.md", result)
 else
     name = first(ARGS)
-    BenchmarkTools.judge(ExponentialFamily, name; judgekwargs = Dict(:time_tolerance => 0.1, :memory_tolerance => 0.05))
-    export_markdown("benchmark_vs_$(name)_result.md", result)
+    result = BenchmarkTools.judge(ExponentialFamily, name; judgekwargs = Dict(:time_tolerance => 0.1, :memory_tolerance => 0.05))
+    export_markdown("./benchmark_logs/benchmark_vs_$(name)_result.md", result)
 end

src/ExponentialFamily.jl

@@ -72,5 +72,6 @@ include("distributions/poisson.jl")
 include("distributions/chi_squared.jl")
 include("distributions/mv_normal_wishart.jl")
 include("distributions/normal_gamma.jl")
+include("distributions/tensor_dirichlet.jl")
 
 end

src/distributions/binomial.jl

@@ -78,6 +78,10 @@ function (::NaturalToMean{Binomial})(tuple_of_η::Tuple{Any}, _)
     return (logistic(η₁),)
 end
 
+function unpack_parameters(::Type{Binomial}, packed, _)
+    return unpack_parameters(Binomial, packed)
+end
+
 function unpack_parameters(::Type{Binomial}, packed)
     return (first(packed),)
 end

src/distributions/categorical.jl

@@ -64,6 +64,10 @@ function (::NaturalToMean{Categorical})(tuple_of_η::Tuple{V}, _) where {V <: Ab
     return (softmax(convert(Vector, η)),)
 end
 
+function unpack_parameters(::Type{Categorical}, packed, _)
+    return unpack_parameters(Categorical, packed)
+end
+
 function unpack_parameters(::Type{Categorical}, packed)
     return (packed,)
 end

src/distributions/laplace.jl

@@ -152,6 +152,10 @@ function (::NaturalToMean{Laplace})(tuple_of_η::Tuple{Any}, _)
     return (-inv(η₁),)
 end
 
+function unpack_parameters(::Type{Laplace}, packed, _)
+    return (first(packed),)
+end
+
 function unpack_parameters(::Type{Laplace}, packed)
     return (first(packed),)
 end

src/distributions/negative_binomial.jl

@@ -92,6 +92,10 @@ function (::NaturalToMean{NegativeBinomial})(tuple_of_η::Tuple{Any}, _)
     return (one(η₁) - exp(η₁),)
 end
 
+function unpack_parameters(::Type{NegativeBinomial}, packed, _)
+    return (first(packed),)
+end
+
 function unpack_parameters(::Type{NegativeBinomial}, packed)
     return (first(packed),)
 end

src/distributions/pareto.jl

@@ -134,6 +134,10 @@ function (::NaturalToMean{Pareto})(tuple_of_η::Tuple{Any}, _)
     return (-η₁ - one(η₁),)
 end
 
+function unpack_parameters(::Type{Pareto}, packed, _)
+    return (first(packed),)
+end
+
 function unpack_parameters(::Type{Pareto}, packed)
     return (first(packed),)
 end