More consistency on arg refactor

Author: SamuelBrand1

src/gen_fn_interface.jl

@@ -373,7 +373,7 @@ function accumulate_param_gradients!(trace)
 end
 
 """
-    (arg_grads, choice_values, choice_grads) = choice_gradients(
+    (arg_grads, choice_values, choice_gradient) = choice_gradients(
         trace, selection=EmptySelection(), retgrad=nothing)
 
 Given a previous trace \$(x, t)\$ (`trace`) and a gradient with respect to the

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
-    (_, choice_values, choice_gradient) = choice_gradients(trace, selection, retval_grad)
-    values = to_array(choice_values, Float64)
-    gradient = to_array(choice_gradient, Float64)
+    (_, value_choices, gradient_choices) = choice_gradients(trace, selection, retval_grad)
+    values = to_array(value_choices, Float64)
+    gradient = to_array(gradient_choices, 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(choice_values, proposed_values)
+    constraints = from_array(value_choices, proposed_values)
     (new_trace, weight, _, discard) = update(trace,
         args, argdiffs, constraints)
     check && check_observations(get_choices(new_trace), observations)
 
     # backward proposal
-    (_, _, backward_choice_gradient) = choice_gradients(new_trace, selection, retval_grad)
-    backward_gradient = to_array(backward_choice_gradient, Float64)
+    (_, _, backward_gradient_choices) = choice_gradients(new_trace, selection, retval_grad)
+    backward_gradient = to_array(backward_gradient_choices, Float64)
     @assert length(backward_gradient) == length(values)
     backward_score = 0.
     backward_mu  = proposed_values + tau * backward_gradient

src/modeling_library/call_at/call_at.jl

@@ -144,11 +144,11 @@ end
 
 function choice_gradients(trace::CallAtTrace, selection::Selection, retval_grad)
     subselection = selection[trace.key]
-    (kernel_input_grads, value_submap, gradient_submap) = choice_gradients(
+    (kernel_input_grads, value_choices_submap, gradient_choices_submap) = choice_gradients(
         trace.subtrace, subselection, retval_grad)
     input_grads = (kernel_input_grads..., nothing)
-    value_choices = CallAtChoiceMap(trace.key, value_submap)
-    gradient_choices = CallAtChoiceMap(trace.key, gradient_submap)
+    value_choices = CallAtChoiceMap(trace.key, value_choices_submap)
+    gradient_choices = CallAtChoiceMap(trace.key, gradient_choices_submap)
     (input_grads, value_choices, gradient_choices)
 end