distributions boilerplate

Author: sbrantq

src/probprog/Distributions.jl

@@ -0,0 +1,35 @@
+abstract type Distribution end
+
+sampler(::Type{D}) where {D<:Distribution} = error("sampler not implemented for $D")
+logpdf_fn(::Type{D}) where {D<:Distribution} = error("logpdf_fn not implemented for $D")
+params(d::Distribution) = error("params not implemented for $(typeof(d))")
+
+struct Normal{Tμ,Tσ,S<:Tuple} <: Distribution
+    μ::Tμ
+    σ::Tσ
+    shape::S
+
+    function Normal{Tμ,Tσ,S}(μ::Tμ, σ::Tσ, shape::S) where {Tμ,Tσ,S<:Tuple}
+        isempty(shape) && throw(ArgumentError("shape cannot be empty"))
+        return new{Tμ,Tσ,S}(μ, σ, shape)
+    end
+end
+
+Normal(μ::Tμ, σ::Tσ, shape::S) where {Tμ,Tσ,S<:Tuple} = Normal{Tμ,Tσ,S}(μ, σ, shape)
+Normal() = Normal(0.0, 1.0, (1,))
+Normal(μ, σ) = Normal(μ, σ, (1,))
+
+params(d::Normal) = (d.μ, d.σ, d.shape)
+
+function _normal_sampler(rng, μ, σ, shape)
+    return μ .+ σ .* randn(rng, shape)
+end
+
+function _normal_logpdf(x, μ, σ, _)
+    z = (x .- μ) ./ σ
+    n = length(x)
+    return -n * log(σ) - n / 2 * log(2π) - sum(z .^ 2) / 2
+end
+
+sampler(::Type{<:Normal}) = _normal_sampler
+logpdf_fn(::Type{<:Normal}) = _normal_logpdf

src/probprog/Modeling.jl

@@ -56,6 +56,12 @@ function sample(
     return traced_result
 end
 
+function sample(
+    rng::AbstractRNG, dist::D; symbol::Symbol=gensym("sample")
+) where {D<:Distribution}
+    return sample(rng, sampler(D), params(dist)...; symbol=symbol, logpdf=logpdf_fn(D))
+end
+
 function untraced_call(rng::AbstractRNG, f::Function, args::Vararg{Any,Nargs}) where {Nargs}
     args_with_rng = (rng, args...)
 

src/probprog/ProbProg.jl

@@ -5,6 +5,7 @@ using ..Reactant:
 using ..Compiler: @jit, @compile
 
 include("Types.jl")
+include("Distributions.jl")
 include("FFI.jl")
 include("Modeling.jl")
 include("Display.jl")
@@ -14,6 +15,9 @@ include("MCMC.jl")
 # Types.
 export ProbProgTrace, Constraint, Selection, Address
 
+# Distributions.
+export Distribution, Normal
+
 # Utility functions.
 export get_choices, select