minor changes

Author: sbrantq

src/ProbProg.jl

@@ -3,10 +3,6 @@ module ProbProg
 using ..Reactant: MLIR, TracedUtils, AbstractConcreteArray
 using Enzyme
 
-function createTrace()
-    return Dict{Symbol,Any}()
-end
-
 function addSampleToTraceLowered(
     trace_ptr_ptr::Ptr{Ptr{Any}},
     symbol_ptr_ptr::Ptr{Ptr{Any}},
@@ -26,6 +22,8 @@ function addSampleToTraceLowered(
         Float32
     elseif datatype_width == 64
         Float64
+    elseif datatype_width == 1
+        Bool
     else
         @ccall printf("Unsupported datatype width: %d\n"::Cstring, datatype_width::Cint)::Cvoid
         return nothing
@@ -268,6 +266,10 @@ end
     return result
 end
 
+function create_trace()
+    return Dict{Symbol,Any}()
+end
+
 function print_trace(trace::Dict{Symbol,Any})
     println("### Probabilistic Program Trace ###")
     for (symbol, sample) in trace

test/probprog/blr.jl

@@ -1,40 +1,45 @@
-using Reactant, Test, Random, StableRNGs, Statistics
+using Reactant, Test, Random
 using Reactant: ProbProg
-using Libdl: Libdl
 
-normal(rng, μ, σ, shape) = μ .+ σ .* randn(rng, shape)
-bernoulli_logit(rng, logit, shape) = rand(rng, shape...) .< (1 ./ (1 .+ exp.(-logit)))
+function normal(rng, μ, σ, shape)
+    return μ .+ σ .* randn(rng, shape)
+end
 
-function blr(seed, N, K)
-    function model(seed, N, K)
-        rng = Random.default_rng()
-        Random.seed!(rng, seed)
+function bernoulli_logit(rng, logit, shape)
+    return rand(rng, shape...) .< (1 ./ (1 .+ exp.(-logit)))
+end
 
-        # α ~ Normal(0, 10, size = 1)
-        α = ProbProg.sample!(normal, rng, 0, 10, (1,); symbol=:α)
+function blr(seed, N, K)
+    rng = Random.default_rng()
+    Random.seed!(rng, seed)
 
-        # β ~ Normal(0, 2.5, size = K)
-        β = ProbProg.sample!(normal, rng, 0, 2.5, (K,); symbol=:β)
+    # α ~ Normal(0, 10, size = 1)
+    α = ProbProg.sample!(normal, rng, 0, 10, (1,); symbol=:α)
 
-        # X ~ Normal(0, 10, size = (N, K))
-        X = ProbProg.sample!(normal, rng, 0, 10, (N, K); symbol=:X) # TODO: double check transpose 
+    # β ~ Normal(0, 2.5, size = K)
+    β = ProbProg.sample!(normal, rng, 0, 2.5, (K,); symbol=:β)
 
-        # μ = α .+ X * β
-        μ = α .+ X * β
+    # X ~ Normal(0, 10, size = (N, K))
+    X = ProbProg.sample!(normal, rng, 0, 10, (N, K); symbol=:X)
 
-        ProbProg.sample!(bernoulli_logit, rng, μ, (N,); symbol=:Y)
+    # μ = α .+ X * β
+    μ = α .+ X * β
 
-        return μ
-    end
+    Y = ProbProg.sample!(bernoulli_logit, rng, μ, (N,); symbol=:Y)
 
-    return ProbProg.simulate!(model, seed, N, K)
+    return Y
 end
 
 @testset "BLR" begin
     N = 5  # number of observations
     K = 3  # number of features
     seed = Reactant.to_rarray(UInt64[1, 4])
 
-    X = ProbProg.getTrace(@jit optimize = :probprog blr(seed, N, K))
-    ProbProg.print_trace(X)
+    trace = ProbProg.create_trace()
+
+    @test size(
+        Array(@jit optimize = :probprog ProbProg.simulate!(blr, seed, N, K; trace))
+    ) == (N,)
+
+    ProbProg.print_trace(trace)
 end

test/probprog/simulate.jl

@@ -22,7 +22,7 @@ using Reactant: ProbProg
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
 
-        trace = ProbProg.createTrace()
+        trace = ProbProg.create_trace()
 
         before = @code_hlo optimize = :no_enzyme simulate_model(trace, seed, μ, σ, shape)
         @test contains(repr(before), "enzyme.simulate")
@@ -40,7 +40,7 @@ using Reactant: ProbProg
         μ = Reactant.ConcreteRNumber(0.0)
         σ = Reactant.ConcreteRNumber(1.0)
 
-        trace = ProbProg.createTrace()
+        trace = ProbProg.create_trace()
 
         result = Array(@jit optimize = :probprog simulate_model(trace, seed, μ, σ, shape))
 
@@ -57,7 +57,7 @@ using Reactant: ProbProg
             return ProbProg.sample!(op, x, y; symbol=:matmul)
         end
 
-        trace = ProbProg.createTrace()
+        trace = ProbProg.create_trace()
         x = reshape(collect(Float64, 1:12), (4, 3))
         y = reshape(collect(Float64, 1:12), (4, 3))
         x_ra = Reactant.to_rarray(x)