Gone for _choices

Renamed 'values_trie' and 'gradient_trie' to 'value_choices' and 'gradient_choices' in HMC and MAP optimization code, as well as corresponding test cases, to improve code readability and clarity.

Author: SamuelBrand1

src/inference/hmc.jl

@@ -88,9 +88,9 @@ function hmc(
 
     # run leapfrog dynamics
     new_trace = trace
-    (_, values_trie, gradient_trie) = choice_gradients(new_trace, selection, retval_grad)
-    values = to_array(values_trie, Float64)
-    gradient = to_array(gradient_trie, Float64)
+    (_, value_choices, gradient_choices) = choice_gradients(new_trace, selection, retval_grad)
+    values = to_array(value_choices, Float64)
+    gradient = to_array(gradient_choices, Float64)
     momenta = sample_momenta(length(values), metric)
     prev_momenta_score = assess_momenta(momenta, metric)
     for step=1:L
@@ -102,10 +102,10 @@ function hmc(
         values += eps * momenta
 
         # get new gradient
-        values_trie = from_array(values_trie, values)
-        (new_trace, _, _) = update(new_trace, args, argdiffs, values_trie)
-        (_, _, gradient_trie) = choice_gradients(new_trace, selection, retval_grad)
-        gradient = to_array(gradient_trie, Float64)
+        value_choices = from_array(value_choices, values)
+        (new_trace, _, _) = update(new_trace, args, argdiffs, value_choices)
+        (_, _, gradient_choices) = choice_gradients(new_trace, selection, retval_grad)
+        gradient = to_array(gradient_choices, Float64)
 
         # half step on momenta
         momenta += (eps / 2) * gradient

src/inference/map_optimize.jl

@@ -11,16 +11,16 @@ function map_optimize(trace, selection::Selection;
     argdiffs = map((_) -> NoChange(), args)
     retval_grad = accepts_output_grad(get_gen_fn(trace)) ? zero(get_retval(trace)) : nothing
 
-    (_, values, gradient) = choice_gradients(trace, selection, retval_grad)
-    values_vec = to_array(values, Float64)
-    gradient_vec = to_array(gradient, Float64)
+    (_, value_choices, gradient_choices) = choice_gradients(trace, selection, retval_grad)
+    values_vec = to_array(value_choices, Float64)
+    gradient_vec = to_array(gradient_choices, Float64)
     step_size = max_step_size
     score = get_score(trace)
     while true
         new_values_vec = values_vec + gradient_vec * step_size
-        values = from_array(values, new_values_vec)
+        value_choices = from_array(value_choices, new_values_vec)
         # TODO discard and weight are not actually needed, there should be a more specialized variant
-        (new_trace, _, _, discard) = update(trace, args, argdiffs, values)
+        (new_trace, _, _, discard) = update(trace, args, argdiffs, value_choices)
         new_score = get_score(new_trace)
         change = new_score - score
         if verbose

test/inference/hmc.jl

@@ -20,9 +20,9 @@
     (new_trace, accepted) = hmc(trace, select(:x))
 
     # For Normal(0,1), grad should be -x
-    (_, values_trie, gradient_trie) = choice_gradients(trace, select(:x), 0)
-    values = to_array(values_trie, Float64)
-    grad = to_array(gradient_trie, Float64)
+    (_, value_choices, gradient_choices) = choice_gradients(trace, select(:x), 0)
+    values = to_array(value_choices, Float64)
+    grad = to_array(gradient_choices, Float64)
     @test values ≈ -grad
 
     # smoke test with vector metric
@@ -58,9 +58,9 @@
     (new_trace, accepted) = hmc(trace, select(:x, :y); metric=metric_vec)
 
     # For each Normal(0,1), grad should be -x
-    (_, values_trie, gradient_trie) = choice_gradients(trace, select(:x, :y), 0)
-    values = to_array(values_trie, Float64)
-    grad = to_array(gradient_trie, Float64)
+    (_, value_choices, gradient_choices) = choice_gradients(trace, select(:x, :y), 0)
+    values = to_array(value_choices, Float64)
+    grad = to_array(gradient_choices, Float64)
     @test values ≈ -grad
 
     # smoke test with Diagonal metric and retval gradient