compiled generate test

Author: sbrantq

test/probprog/generate.jl

@@ -39,4 +39,38 @@ end
             normal_logpdf(trace.choices[:t][1], constraint[:s][1], 1.0, shape)
         @test weight ≈ expected_weight atol = 1e-6
     end
+
+    @testset "compiled" begin
+        shape = (10,)
+        seed = Reactant.to_rarray(UInt64[1, 4])
+        μ = Reactant.ConcreteRNumber(0.0)
+        σ = Reactant.ConcreteRNumber(1.0)
+
+        constraint1 = Dict{Symbol,Any}(:s => (fill(0.1, shape),))
+
+        constrained_symbols = collect(keys(constraint1)) # This doesn't change
+
+        constraint_ptr1 = Reactant.ConcreteRNumber(
+            reinterpret(UInt64, pointer_from_objref(constraint1))
+        )
+
+        wrapper_fn(constraint_ptr, seed, μ, σ) = ProbProg.generate_internal(
+            model, seed, μ, σ, shape; constraint_ptr, constrained_symbols
+        )
+
+        compiled_fn = @compile optimize = :probprog wrapper_fn(constraint_ptr1, seed, μ, σ)
+
+        trace1, weight = compiled_fn(constraint_ptr1, seed, μ, σ)
+        trace1 = unsafe_pointer_to_objref(Ptr{Any}(Array(trace1)[1]))
+
+        constraint2 = Dict{Symbol,Any}(:s => (fill(0.2, shape),))
+        constraint_ptr2 = Reactant.ConcreteRNumber(
+            reinterpret(UInt64, pointer_from_objref(constraint2))
+        )
+
+        trace2, _ = compiled_fn(constraint_ptr2, seed, μ, σ)
+        trace2 = unsafe_pointer_to_objref(Ptr{Any}(Array(trace2)[1]))
+
+        @test trace1.choices[:s] != trace2.choices[:s]
+    end
 end