Merge pull request #497 from probcomp/20230109-ztangent-fix_dist_dsl_grads

Fix `logpdf_grad` errors in `@dist` DSL.

Author: alex-lew

src/modeling_library/dist_dsl/dist_dsl.jl

@@ -26,13 +26,73 @@ all_indices(arg::SimpleArg) = [arg.i]
 all_indices(arg::TransformedArg) = vcat([all_indices(a) for a in arg.f_args]...)
 
 # Evaluate user-facing args to concrete values passed to the base distribution
-eval_arg(x::Any, args) = x
-eval_arg(x::SimpleArg, args) = typecheck_arg(x, args[x.i])
-eval_arg(x::TransformedArg, args) =
-    x.arg_passer(x.orig_f, [eval_arg(a, args) for a in x.f_args]...)
+eval_arg(base_arg::Any, args) = base_arg
+eval_arg(base_arg::SimpleArg, args) = typecheck_arg(base_arg, args[base_arg.i])
+eval_arg(base_arg::TransformedArg, args) =
+    base_arg.arg_passer(base_arg.orig_f, (eval_arg(a, args) for a in base_arg.f_args)...)
+
+# Evaluate gradients of base distribution args with respect to user-facing args
+function eval_arg_gradient(base_arg::Any, base_type::Type, args)
+    grads = map(enumerate(args)) do (i, arg)
+        if arg isa Real || arg isa AbstractArray && eltype(arg) <: Real
+            zero(arg) # Base arg is always constant with respect to input args
+        else
+            nothing
+        end
+    end
+    return grads
+end
+
+function eval_arg_gradient(base_arg::SimpleArg{T}, base_type::Type, args) where {T}
+    grads = map(enumerate(args)) do (i, arg)
+        if arg isa Real # Base arg is either equal to or unaffected by input arg
+            i == base_arg.i ? one(arg) : zero(arg)
+        elseif arg isa AbstractArray && eltype(arg) <: Real
+            N, V = length(arg), eltype(arg)
+            i == base_arg.i ? Matrix{V}(LinearAlgebra.I, N, N) : zeros(V, N, N) 
+        else
+            nothing
+        end
+    end
+    return grads
+end
+
+# Compute gradients when base arg is a scalar type
+function eval_arg_gradient(base_arg::TransformedArg, base_type::Type{<:Real}, args)
+    splice_arg(arg, i) = [args[1:i-1]..., arg, args[i+1:end]...]
+    per_arg_eval(arg, i) = eval_arg(base_arg, splice_arg(arg, i))
+    grads = map(enumerate(args)) do (i, arg)
+        if arg isa Real
+            ReverseDiff.gradient(a -> per_arg_eval(a, i), [arg])[1]
+        elseif arg isa AbstractArray && eltype(arg) <: Real
+            ReverseDiff.gradient(a -> per_arg_eval(a, i), arg)
+        else
+            nothing
+        end
+    end
+    return grads
+end
+
+# Compute Jacobians when base arg is an array type
+function eval_arg_gradient(base_arg::TransformedArg, base_type::Type{<:AbstractArray{<:Real}}, args)
+    splice_arg(arg, i) = [args[1:i-1]..., arg, args[i+1:end]...]
+    per_arg_eval(arg, i) = eval_arg(base_arg, splice_arg(arg, i))
+    grads = map(enumerate(args)) do (i, arg)
+        if arg isa Real
+            ReverseDiff.jacobian(a -> per_arg_eval(a, i), [arg])
+        elseif arg isa AbstractArray && eltype(arg) <: Real
+            ReverseDiff.jacobian(a -> per_arg_eval(a, i), arg)
+        else
+            nothing
+        end
+    end
+    return grads
+end
 
 # Type of SimpleArg must match arg, otherwise a MethodError will be thrown
-typecheck_arg(x::SimpleArg{T}, arg::T) where {T} = arg
+typecheck_arg(base_arg::SimpleArg{T}, arg::T) where {T} = arg
+typecheck_arg(base_arg::SimpleArg{T}, arg::ReverseDiff.TrackedReal{T}) where {T <: Real} = arg
+typecheck_arg(base_arg::SimpleArg{T}, arg::ReverseDiff.TrackedArray{V, D, N, T}) where {V, D, N, T} = arg
 
 # DistWithArgs
 struct DistWithArgs{T}
@@ -71,22 +131,32 @@ function logpdf_grad(d::CompiledDistWithArgs{T}, x::T, args...) where T
     concrete_args = [eval_arg(arg, args) for arg in d.arglist]
     base_has_arg_grads = has_argument_grads(d.base)
     base_grads = logpdf_grad(d.base, x, concrete_args...)
-
-    base_arg_grads = [g for (i, g) in enumerate(base_grads[2:end])
-                      if base_has_arg_grads[i]]
-    argvec = collect(args)
-    eval_arg_grads = hcat([ReverseDiff.gradient(xs -> eval_arg(arg, xs), argvec)
-        for (i, arg) in enumerate(d.arglist) if base_has_arg_grads[i]]...)
-
-    retval = [base_grads[1]]
-    for i in 1:d.n_args
-        if self_has_arg_grads[i]
-            push!(retval, eval_arg_grads[i,:]' * base_arg_grads)
-        else
-            push!(retval, nothing)
+    base_arg_grads = base_grads[2:end]
+
+    # Set gradient with respect to output
+    self_output_grad = base_grads[1] 
+
+    # Backpropagate gradients from base arguments to arguments
+    self_arg_grads = [self_has_arg_grads[i] ? zero(arg) : nothing
+                      for (i, arg) in enumerate(args)]
+
+    for (i, base_arg) in enumerate(d.arglist)
+        base_has_arg_grads[i] || continue
+        base_grad = base_arg_grads[i]
+        base_arg_type = typeof(concrete_args[i])
+        eval_arg_grad = eval_arg_gradient(base_arg, base_arg_type, args) 
+        for (j, g) in enumerate(eval_arg_grad)
+            (isnothing(g) || !self_has_arg_grads[j]) && continue
+            if base_grad isa AbstractArray
+                increment = reshape(g' * vec(base_grad), size(self_arg_grads[j]))
+            else
+                increment = g * base_grad
+            end
+            self_arg_grads[j] = self_arg_grads[j] .+ increment
         end
     end
-    retval
+
+    return (self_output_grad, self_arg_grads...)
 end
 
 function random(d::CompiledDistWithArgs{T}, args...)::T where T

src/modeling_library/dist_dsl/relabeled_distribution.jl

@@ -19,13 +19,13 @@ function logpdf(d::WithLabelArg{T, U}, x::T, collection, base_args...) where {T,
 end
 
 function logpdf_grad(d::WithLabelArg{T, U}, x::T, collection, base_args...) where {T, U}
-    base_arg_grads = fill(nothing, length(base_args))
+    base_arg_grads = Vector{Any}(nothing, length(base_args))
 
     for p in pairs(collection)
         (index, item) = (p.first, p.second)
         if item == x
             new_grads = logpdf_grad(d.base, index, base_args...)
-            for (arg_idx, grad) in enumerate(new_grads)
+            for (arg_idx, grad) in enumerate(new_grads[2:end])
                 if base_arg_grads[arg_idx] === nothing
                     base_arg_grads[arg_idx] = grad
                 elseif grad !== nothing
@@ -73,13 +73,13 @@ function logpdf(d::RelabeledDistribution{T, U}, x::T, base_args...) where {T, U}
 end
 
 function logpdf_grad(d::RelabeledDistribution{T, U}, x::T, base_args...) where {T, U}
-    base_arg_grads = fill(nothing, length(base_args))
+    base_arg_grads = Vector{Any}(nothing, length(base_args))
 
     for p in pairs(d.collection)
         (index, item) = (p.first, p.second)
         if item == x
             new_grads = logpdf_grad(d.base, index, base_args...)
-            for (arg_idx, grad) in enumerate(new_grads)
+            for (arg_idx, grad) in enumerate(new_grads[2:end])
                 if base_arg_grads[arg_idx] === nothing
                     base_arg_grads[arg_idx] = grad
                 elseif grad !== nothing

src/modeling_library/dist_dsl/transformed_distribution.jl

@@ -41,14 +41,14 @@ function logpdf_grad(d::TransformedDistribution{T, U}, x::T, args...) where {T,
 
     if is_discrete(d.base) || !has_output_grad(d.base)
         # TODO: should this be nothing or 0?
-        [base_grad[1], fill(nothing, d.nArgs)..., base_grad[2:end]...]
+        return (base_grad[1], fill(nothing, d.nArgs)..., base_grad[2:end]...)
     else
         transformation_grad = d.backward_grad(x, args[1:d.nArgs]...)
         correction_grad = ReverseDiff.gradient(v -> logpdf_correction(d, v[1], v[2:end]), [x, args[1:d.nArgs]...])
         # TODO: Will this sort of thing work if the arguments w.r.t. which we are taking
         # gradients are themselves vector-valued?
         full_grad = (transformation_grad .* base_grad[1]) .+ correction_grad
-        [full_grad..., base_grad[2:end]...]
+        return (full_grad..., base_grad[2:end]...)
     end
 end
 
@@ -62,8 +62,8 @@ end
 
 function has_argument_grads(d::TransformedDistribution{T, U}) where {T, U}
     if is_discrete(d.base) || !has_output_grad(d.base)
-        [fill(false, d.nArgs)..., has_argument_grads(d.base)...]
+        (fill(false, d.nArgs)..., has_argument_grads(d.base)...)
     else
-        [fill(true, d.nArgs)..., has_argument_grads(d.base)...]
+        (fill(true, d.nArgs)..., has_argument_grads(d.base)...)
     end
 end

test/modeling_library/dist_dsl.jl

@@ -2,18 +2,51 @@
   # Test transformed distributions
   @dist f(x) = exp(normal(x, 0.001))
   @test isapprox(1, f(0); atol = 5)
+  @test isapprox(logpdf(f, 1., 0.), logpdf(normal, 0., 0., 0.001))
+
+  # Test gradients of transformed distributions
+  @dist shifted_normal(mu, sigma) = normal(mu, sigma) + 1.
+  @test logpdf(shifted_normal, 1., 0., 1.) == logpdf(normal, 0., 0., 1.)
+  @test logpdf_grad(shifted_normal, 0., 0., 1.) == logpdf_grad(normal, -1., 0., 1.)
+
+  # Test gradients of transformed distributions with no parameters
+  @dist shifted_std_normal() = normal(0., 1.) + 1.
+  @test logpdf(shifted_std_normal, 1.) == logpdf(normal, 0., 0., 1.)
+  @test logpdf_grad(shifted_std_normal, 0.) == (logpdf_grad(normal, -1., 0., 1.)[1],)
+
+  # Test gradients of multivariate distributions
+  @dist vec_normal(mu, sigma) = broadcasted_normal(broadcast(+, mu, 1.0), broadcast(*, sigma, 2.0))
+  @test logpdf(vec_normal, zeros(2), zeros(2), ones(2)) ==
+        logpdf(broadcasted_normal, zeros(2), ones(2), 2 .* ones(2))
+  transformed_grads = logpdf_grad(vec_normal, zeros(2), zeros(2), ones(2))
+  orig_grads = logpdf_grad(broadcasted_normal, zeros(2), ones(2), 2 .* ones(2)) 
+  @test transformed_grads[1] == orig_grads[1]
+  @test transformed_grads[2] == orig_grads[2]
+  @test transformed_grads[3] == 2.0 * orig_grads[3]
+      
+  # Test gradients of multivariate distributions with multi-dimensional arguments
+  @dist transformed_mvnormal(mu, sigma) = mvnormal(broadcast(+, mu, 1.0), broadcast(*, sigma, 2.0))
+  @test logpdf(transformed_mvnormal, zeros(2), zeros(2), ones(2, 2)) ==
+        logpdf(mvnormal, zeros(2), ones(2), 2 .* ones(2, 2))
+  transformed_grads = logpdf_grad(transformed_mvnormal, zeros(2), zeros(2), ones(2, 2))
+  orig_grads = logpdf_grad(mvnormal, zeros(2), ones(2), 2 .* ones(2, 2)) 
+  @test transformed_grads[1] == orig_grads[1]
+  @test transformed_grads[2] == orig_grads[2]
+  @test transformed_grads[3] == 2.0 * orig_grads[3]
 
   # Test relabeled distributions with labels provided as an Array
   @dist labeled_cat(labels, probs) = labels[categorical(probs)]
   @test labeled_cat([:a, :b], [0., 1.]) == :b
   @test isapprox(logpdf(labeled_cat, :b, [:a, :b], [0.5, 0.5]), log(0.5))
+  @test logpdf_grad(labeled_cat, :b, [:a, :b], [0.5, 0.5]) == (nothing, logpdf_grad(categorical, 2, [0.5, 0.5])...)
   @test logpdf(labeled_cat, :c, [:a, :b], [0.5, 0.5]) == -Inf
 
   # Test relabeled distributions with labels provided in a Dict
   dict = Dict(1 => :a, 2 => :b)
   @dist dict_cat(probs) = dict[categorical(probs)]
   @test dict_cat([0., 1.]) == :b
   @test isapprox(logpdf(dict_cat, :b, [0.5, 0.5]), log(0.5))
+  @test logpdf_grad(dict_cat, :b, [0.5, 0.5]) == logpdf_grad(categorical, 2, [0.5, 0.5])
   @test logpdf(dict_cat, :c, [0.5, 0.5]) == -Inf
 
   # Test relabeled distributions with Enum labels
@@ -22,11 +55,13 @@
   @test enum_cat([0., 1.]) == orange
   @test isapprox(logpdf(enum_cat, orange, [0.5, 0.5]), log(0.5))
   @test logpdf(enum_cat, orange, [1.0]) == -Inf
+  @test logpdf_grad(enum_cat, orange, [0.5, 0.5]) == logpdf_grad(categorical, 2, [0.5, 0.5])
 
   # Regression test for https://github.com/probcomp/Gen/issues/253
   @dist real_minus_uniform(a, b) = 1 - Gen.uniform(a, b)
   @test real_minus_uniform(1, 2) < 0
   @test logpdf(real_minus_uniform, -0.5, 1, 2) == 0.0
+  @test logpdf_grad(real_minus_uniform, -0.5, 1, 2) == logpdf_grad(uniform, 1.5, 1, 2)
 end
 
 # User-defined type for testing purposes
@@ -40,13 +75,15 @@ end
   @test symbol_cat([:a, :b], [0., 1.]) == :b
   @test_throws MethodError symbol_cat(["a", "b"], [0., 1.])
   @test logpdf(symbol_cat, :c, [:a, :b], [0.5, 0.5]) == -Inf
+  @test logpdf_grad(symbol_cat, :b, [:a, :b], [0.5, 0.5]) == (nothing, logpdf_grad(categorical, 2, [0.5, 0.5])...)
   @test_throws MethodError logpdf(symbol_cat, "c", [:a, :b], [0.5, 0.5])
 
   # Test typed parameters
   @dist int_bounded_uniform(low::Int, high::Int) = uniform(low, high)
   @test 0.0 <= int_bounded_uniform(0, 1) <= 1
   @test_throws MethodError int_bounded_uniform(-0.5, 0.5)
   @test logpdf(int_bounded_uniform, 0.5, 0, 1) == 0
+  @test logpdf_grad(int_bounded_uniform, 0.5, 0, 1) == logpdf_grad(uniform, 0.5, 0, 1)
   @test_throws MethodError logpdf(int_bounded_uniform, 0.0, -0.5, 0.5)
 
   # Test relabeled distributions with user-defined types