Merge pull request #448 from probcomp/mrb/unfold_gradient_default_args

Fix Unfold propagating gradients to default parameters.

Author: alex-lew

.gitignore

@@ -1,3 +1,4 @@
+Manifest.toml
 *.pdf
 *.png
 *.jld
@@ -7,4 +8,4 @@
 *.log
 docs/build/
 docs/site/
-.DS_Store
+.DS_Store

src/modeling_library/unfold/backprop.jl

@@ -4,6 +4,9 @@
 @inline fold_sum(::Nothing, a::A) where A = a
 @inline fold_sum(a::A, b::A) where A = a + b
 
+@inline _sum(::Nothing) = nothing
+@inline _sum(x::Vector) = sum(x)
+
 function choice_gradients(trace::VectorTrace{UnfoldType,T,U}, selection::Selection, retval_grad) where {T,U}
     kernel_has_grads = has_argument_grads(trace.gen_fn.kernel)
     if kernel_has_grads[1]
@@ -44,7 +47,7 @@ function choice_gradients(trace::VectorTrace{UnfoldType,T,U}, selection::Selecti
             end
         end
     end
-    ((nothing, kernel_arg_grads[2], params_grad...), value_choices, gradient_choices)
+    ((nothing, kernel_arg_grads[2], map(_sum, params_grad)...), value_choices, gradient_choices)
 end
 
 function accumulate_param_gradients!(trace::VectorTrace{UnfoldType,T,U}, retval_grad) where {T,U}
@@ -82,5 +85,5 @@ function accumulate_param_gradients!(trace::VectorTrace{UnfoldType,T,U}, retval_
             end
         end
     end
-    (nothing, kernel_arg_grads[2], params_grad...)
+    (nothing, kernel_arg_grads[2], map(_sum, params_grad)...)
 end

test/modeling_library/unfold.jl

@@ -506,10 +506,13 @@ foo = Unfold(kernel)
 
         zero_param_grad!(kernel, :std)
         input_grads = accumulate_param_gradients!(trace, nothing)
+        expected_xs_grad = logpdf_grad(normal, x1, x_init * alpha + beta, std)[2] * x_init + logpdf_grad(normal, x2, x1 * alpha + beta, std)[2] * x1
+        expected_ys_grad = logpdf_grad(normal, x1, x_init * alpha + beta, std)[2] + logpdf_grad(normal, x2, x1 * alpha + beta, std)[2]
+
         @test input_grads[1] == nothing # length
         @test input_grads[2] == nothing # inital state
-        #@test isapprox(input_grads[3], expected_xs_grad) # alpha
-        #@test isapprox(input_grads[4], expected_ys_grad) # beta
+        @test isapprox(input_grads[3], expected_xs_grad) # alpha
+        @test isapprox(input_grads[4], expected_ys_grad) # beta
         expected_std_grad = (logpdf_grad(normal, x1, x_init * alpha + beta, std)[3]
                              + logpdf_grad(normal, x2, x1 * alpha + beta, std)[3])
         @test isapprox(get_param_grad(kernel, :std), expected_std_grad)