update legacy inference API

Author: sbrantq

src/probprog/Inference.jl

@@ -10,17 +10,17 @@ function metropolis_hastings(
         error("MH requires a trace with fn and rng recorded (use generate to create trace)")
     end
 
-    constraints = Dict{Symbol,Any}()
-    constrained_symbols = Set{Symbol}()
-
+    constraint_pairs = Pair{Symbol,Any}[]
     for (sym, val) in trace.choices
         if !(sym in sel)
-            constraints[sym] = val
-            push!(constrained_symbols, sym)
+            push!(constraint_pairs, sym => val)
         end
     end
+    constraint = Constraint(constraint_pairs...)
+
+    constrained_addresses = extract_addresses(constraint)
 
-    cache_key = (typeof(trace.fn), constrained_symbols)
+    cache_key = (typeof(trace.fn), constrained_addresses)
 
     compiled_fn = nothing
     if compiled_cache !== nothing
@@ -30,12 +30,12 @@ function metropolis_hastings(
     if compiled_fn === nothing
         function wrapper_fn(rng, constraint_ptr, args...)
             return generate_internal(
-                rng, trace.fn, args...; constraint_ptr, constrained_symbols
+                rng, trace.fn, args...; constraint_ptr, constrained_addresses
             )
         end
 
         constraint_ptr = ConcreteRNumber(
-            reinterpret(UInt64, pointer_from_objref(constraints))
+            reinterpret(UInt64, pointer_from_objref(constraint))
         )
 
         compiled_fn = @compile optimize = :probprog wrapper_fn(
@@ -47,7 +47,7 @@ function metropolis_hastings(
         end
     end
 
-    constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraints)))
+    constraint_ptr = ConcreteRNumber(reinterpret(UInt64, pointer_from_objref(constraint)))
 
     old_gc_state = GC.enable(false)
     new_trace_ptr = nothing

src/probprog/Types.jl

@@ -70,7 +70,7 @@ Base.get(c::Constraint, k::Address, default) = get(c.dict, k, default)
 extract_addresses(constraint::Constraint) = Set(keys(constraint))
 
 const Selection = Set{Symbol}
-const CompiledFnCache = Dict{Tuple{Type,Set{Symbol}},Any}
+const CompiledFnCache = Dict{Tuple{Type,Set{Address}},Any}
 
 const _probprog_ref_lock = ReentrantLock()
 const _probprog_refs = IdDict()

test/probprog/linear_regression.jl

@@ -33,13 +33,12 @@ end
 function my_inference_program(xs, ys, num_iters)
     xs_r = Reactant.to_rarray(xs)
 
-    constraint = ProbProg.choicemap()
-    constraint[:ys] = [ys]
+    observations = ProbProg.Constraint(:ys => (ys,))
 
     seed = Reactant.to_rarray(UInt64[1, 4])
     rng = ReactantRNG(seed)
 
-    trace, _ = ProbProg.generate(rng, my_model, xs_r; constraint)
+    trace, _ = ProbProg.generate(rng, my_model, xs_r; constraint=observations)
 
     trace = ProbProg.with_compiled_cache() do cache
         local t = trace