Rename value_trie and gradient_trie to choice_value and choice_gradient

Refactored code and tests to consistently use 'choice_value' and 'choice_gradient' instead of 'value_trie' and 'gradient_trie' for clarity and alignment with choicemap terminology. Updated variable names, function signatures, and related comments across inference, static IR, and test files.

Author: SamuelBrand1

.vscode/settings.json

@@ -0,0 +1 @@
+{}

src/inference/mala.jl

@@ -17,9 +17,9 @@ function mala(
     retval_grad = accepts_output_grad(get_gen_fn(trace)) ? zero(get_retval(trace)) : nothing
 
     # forward proposal
-    (_, values_trie, gradient_trie) = choice_gradients(trace, selection, retval_grad)
-    values = to_array(values_trie, Float64)
-    gradient = to_array(gradient_trie, Float64)
+    (_, choice_values, choice_gradient) = choice_gradients(trace, selection, retval_grad)
+    values = to_array(choice_values, Float64)
+    gradient = to_array(choice_gradient, Float64)
     forward_mu = values + tau * gradient
     forward_score = 0.
     proposed_values = Vector{Float64}(undef, length(values))
@@ -29,14 +29,14 @@ function mala(
     end
 
     # evaluate model weight
-    constraints = from_array(values_trie, proposed_values)
+    constraints = from_array(choice_values, proposed_values)
     (new_trace, weight, _, discard) = update(trace,
         args, argdiffs, constraints)
     check && check_observations(get_choices(new_trace), observations)
 
     # backward proposal
-    (_, _, backward_gradient_trie) = choice_gradients(new_trace, selection, retval_grad)
-    backward_gradient = to_array(backward_gradient_trie, Float64)
+    (_, _, backward_choice_gradient) = choice_gradients(new_trace, selection, retval_grad)
+    backward_gradient = to_array(backward_choice_gradient, Float64)
     @assert length(backward_gradient) == length(values)
     backward_score = 0.
     backward_mu  = proposed_values + tau * backward_gradient

src/static_ir/backprop.jl

@@ -4,11 +4,11 @@ struct BackpropParamsMode end
 const gradient_prefix = gensym("gradient")
 gradient_var(node::StaticIRNode) = Symbol("$(gradient_prefix)_$(node.name)")
 
-const value_trie_prefix = gensym("value_trie")
-value_trie_var(node::GenerativeFunctionCallNode) = Symbol("$(value_trie_prefix)_$(node.addr)")
+const choice_value_prefix = gensym("choice_value")
+choice_value_var(node::GenerativeFunctionCallNode) = Symbol("$(choice_value_prefix)_$(node.addr)")
 
-const gradient_trie_prefix = gensym("gradient_trie")
-gradient_trie_var(node::GenerativeFunctionCallNode) = Symbol("$(gradient_trie_prefix)_$(node.addr)")
+const choice_gradient_prefix = gensym("choice_gradient")
+choice_gradient_var(node::GenerativeFunctionCallNode) = Symbol("$(choice_gradient_prefix)_$(node.addr)")
 
 const tape_prefix = gensym("tape")
 tape_var(node::JuliaNode) = Symbol("$(tape_prefix)_$(node.name)")
@@ -128,7 +128,7 @@ function fwd_codegen!(stmts, fwd_marked, back_marked, node::JuliaNode)
         # we need the value for initializing gradient to zero (to get the type
         # and e.g. shape), and for reference by other nodes during
         # back_codegen! we could be more selective about which JuliaNodes need
-        # to be evalutaed, that is a performance optimization for the future
+        # to be evaluated, that is a performance optimization for the future
         args = map((input_node) -> input_node.name, node.inputs)
         push!(stmts, :($(node.name) = $(QuoteNode(node.fn))($(args...))))
     end
@@ -275,8 +275,8 @@ function back_codegen!(stmts, ir, selected_calls, fwd_marked, back_marked,
 
     if node in fwd_marked
         input_grads = gensym("call_input_grads")
-        value_trie = value_trie_var(node)
-        gradient_trie = gradient_trie_var(node)
+        choice_value = choice_value_var(node)
+        choice_gradient = choice_gradient_var(node)
         subtrace_fieldname = get_subtrace_fieldname(node)
         call_selection = gensym("call_selection")
         if node in selected_calls
@@ -285,7 +285,7 @@ function back_codegen!(stmts, ir, selected_calls, fwd_marked, back_marked,
             push!(stmts, :($call_selection = EmptySelection()))
         end
         retval_grad = node in back_marked ? gradient_var(node) : :(nothing)
-        push!(stmts, :(($input_grads, $value_trie, $gradient_trie) = choice_gradients(
+        push!(stmts, :(($input_grads, $choice_value, $choice_gradient) = choice_gradients(
             trace.$subtrace_fieldname, $call_selection, $retval_grad)))
     end
 
@@ -297,7 +297,7 @@ function back_codegen!(stmts, ir, selected_calls, fwd_marked, back_marked,
         end
     end
 
-    # NOTE: the value_trie and gradient_trie are dealt with later
+    # NOTE: the choice_value and choice_gradient are dealt with later
 end
 
 function back_codegen!(stmts, ir, selected_calls, fwd_marked, back_marked,
@@ -325,9 +325,9 @@ function back_codegen!(stmts, ir, selected_calls, fwd_marked, back_marked,
     end
 end
 
-function generate_value_gradient_trie(selected_choices::Set{RandomChoiceNode},
+function generate_value_choice_gradient(selected_choices::Set{RandomChoiceNode},
                                       selected_calls::Set{GenerativeFunctionCallNode},
-                                      value_trie::Symbol, gradient_trie::Symbol)
+                                      choice_value::Symbol, choice_gradient::Symbol)
     selected_choices_vec = collect(selected_choices)
     quoted_leaf_keys = map((node) -> QuoteNode(node.addr), selected_choices_vec)
     leaf_values = map((node) -> :(trace.$(get_value_fieldname(node))), selected_choices_vec)
@@ -337,12 +337,12 @@ function generate_value_gradient_trie(selected_choices::Set{RandomChoiceNode},
     quoted_internal_keys = map((node) -> QuoteNode(node.addr), selected_calls_vec)
     internal_values = map((node) -> :(get_choices(trace.$(get_subtrace_fieldname(node)))),
                           selected_calls_vec)
-    internal_gradients = map((node) -> gradient_trie_var(node), selected_calls_vec)
+    internal_gradients = map((node) -> choice_gradient_var(node), selected_calls_vec)
     quote
-        $value_trie = StaticChoiceMap(
+        $choice_value = StaticChoiceMap(
             NamedTuple{($(quoted_leaf_keys...),)}(($(leaf_values...),)),
             NamedTuple{($(quoted_internal_keys...),)}(($(internal_values...),)))
-        $gradient_trie = StaticChoiceMap(
+        $choice_gradient = StaticChoiceMap(
             NamedTuple{($(quoted_leaf_keys...),)}(($(leaf_gradients...),)),
             NamedTuple{($(quoted_internal_keys...),)}(($(internal_gradients...),)))
     end
@@ -429,18 +429,18 @@ function codegen_choice_gradients(trace_type::Type{T}, selection_type::Type,
         back_codegen!(stmts, ir, selected_calls, fwd_marked, back_marked, node, BackpropTraceMode())
     end
 
-    # assemble value_trie and gradient_trie
-    value_trie = gensym("value_trie")
-    gradient_trie = gensym("gradient_trie")
-    push!(stmts, generate_value_gradient_trie(selected_choices, selected_calls,
-        value_trie, gradient_trie))
+    # assemble choice_value and choice_gradient
+    choice_value = gensym("choice_value")
+    choice_gradient = gensym("choice_gradient")
+    push!(stmts, generate_value_choice_gradient(selected_choices, selected_calls,
+        choice_value, choice_gradient))
 
     # gradients with respect to inputs
     arg_grad = (node::ArgumentNode) -> node.compute_grad ? gradient_var(node) : :(nothing)
     input_grads = Expr(:tuple, map(arg_grad, ir.arg_nodes)...)
 
     # return values
-    push!(stmts, :(return ($input_grads, $value_trie, $gradient_trie)))
+    push!(stmts, :(return ($input_grads, $choice_value, $choice_gradient)))
 
     Expr(:block, stmts...)
 end

test/dsl/dynamic_dsl.jl

@@ -345,15 +345,15 @@ end
     # check input gradient
     @test isapprox(mu_a_grad, finite_diff(f, (mu_a, a, b, z, out), 1, dx))
 
-    # check value trie
+    # check value from choicemap
     @test get_value(choices, :a) == a
     @test get_value(choices, :out) == out
     @test get_value(choices, :bar => :z) == z
     @test !has_value(choices, :b) # was not selected
     @test length(collect(get_submaps_shallow(choices))) == 1
     @test length(collect(get_values_shallow(choices))) == 2
 
-    # check gradient trie
+    # check gradient from choicemap
     @test length(collect(get_submaps_shallow(gradients))) == 1
     @test length(collect(get_values_shallow(gradients))) == 2
     @test !has_value(gradients, :b) # was not selected

test/static_ir/static_ir.jl

@@ -351,26 +351,26 @@ end
     selection = select(:bar => :z, :a, :out)
     selection = StaticSelection(selection)
     retval_grad = 2.
-    ((mu_a_grad,), value_trie, gradient_trie) = choice_gradients(trace, selection, retval_grad)
+    ((mu_a_grad,), choice_value, choice_gradient) = choice_gradients(trace, selection, retval_grad)
 
     # check input gradient
     @test isapprox(mu_a_grad, finite_diff(f, (mu_a, theta, a, b, z, out), 1, dx))
 
-    # check value trie
-    @test get_value(value_trie, :a) == a
-    @test get_value(value_trie, :out) == out
-    @test get_value(value_trie, :bar => :z) == z
-    @test !has_value(value_trie, :b) # was not selected
-    @test length(get_submaps_shallow(value_trie)) == 1
-    @test length(get_values_shallow(value_trie)) == 2
-
-    # check gradient trie
-    @test length(get_submaps_shallow(gradient_trie)) == 1
-    @test length(get_values_shallow(gradient_trie)) == 2
-    @test !has_value(gradient_trie, :b) # was not selected
-    @test isapprox(get_value(gradient_trie, :a), finite_diff(f, (mu_a, theta, a, b, z, out), 3, dx))
-    @test isapprox(get_value(gradient_trie, :out), finite_diff(f, (mu_a, theta, a, b, z, out), 6, dx))
-    @test isapprox(get_value(gradient_trie, :bar => :z), finite_diff(f, (mu_a, theta, a, b, z, out), 5, dx))
+    # check value from choice map
+    @test get_value(choice_value, :a) == a
+    @test get_value(choice_value, :out) == out
+    @test get_value(choice_value, :bar => :z) == z
+    @test !has_value(choice_value, :b) # was not selected
+    @test length(get_submaps_shallow(choice_value)) == 1
+    @test length(get_values_shallow(choice_value)) == 2
+
+    # check gradient from choice map
+    @test length(get_submaps_shallow(choice_gradient)) == 1
+    @test length(get_values_shallow(choice_gradient)) == 2
+    @test !has_value(choice_gradient, :b) # was not selected
+    @test isapprox(get_value(choice_gradient, :a), finite_diff(f, (mu_a, theta, a, b, z, out), 3, dx))
+    @test isapprox(get_value(choice_gradient, :out), finite_diff(f, (mu_a, theta, a, b, z, out), 6, dx))
+    @test isapprox(get_value(choice_gradient, :bar => :z), finite_diff(f, (mu_a, theta, a, b, z, out), 5, dx))
 
     # reset the trainable parameter gradient
     zero_param_grad!(foo, :theta)