Added default impls for filldist and arraydist (#264)

* added default impls for `filldist` and `arraydist` * Apply suggestions from code review Co-authored-by: David Widmann <[email protected]> * Update filldist.jl (#265) --------- Co-authored-by: David Widmann <[email protected]>
TuringLang · Apr 17, 2024 · e3f8a67 · e3f8a67 · yebai · Apr 17, 2024
1 parent b8c3d82
commit e3f8a67
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "DistributionsAD"
 uuid = "ced4e74d-a319-5a8a-b0ac-84af2272839c"
-version = "0.6.54"
+version = "0.6.55"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/arraydist.jl b/src/arraydist.jl
@@ -1,3 +1,10 @@
+"""
+    arraydist(dists)
+
+Create a distribution from an array of distributions.
+"""
+arraydist(dists::AbstractArray{<:Distribution}) = product_distribution(dists)
+
 # Univariate
 
 const VectorOfUnivariate = Distributions.Product

diff --git a/src/filldist.jl b/src/filldist.jl
@@ -1,5 +1,14 @@
+# Default implementation just defers to Distributions.jl.
+"""
+    filldist(d::Distribution, ns...)
+
+Create a product distribution using `FillArrays.Fill` as the array type.
+"""
+filldist(d::Distribution, n1::Int, ns::Int...) = product_distribution(Fill(d, n1, ns...))
+
 # Univariate
 
+# TODO: Do we even need these? Probably should benchmark to be sure.
 const FillVectorOfUnivariate{
     S <: ValueSupport,
     T <: UnivariateDistribution{S},
@@ -59,7 +68,7 @@ const FillMatrixOfUnivariate{
     Tdists <: Fill{T, 2},
 } = MatrixOfUnivariate{S, T, Tdists}
 
-function filldist(dist::UnivariateDistribution, N1::Integer, N2::Integer)
+function filldist(dist::UnivariateDistribution, N1::Int, N2::Int)
     return MatrixOfUnivariate(Fill(dist, N1, N2))
 end
 function Distributions._logpdf(dist::FillMatrixOfUnivariate, x::AbstractMatrix{<:Real})