Fix logpdf_grad for BroadcastedNormal.

[MrVPlusOne]

As requested by #432, this PR fixes the gradient of BroadcastedNormal for the non-scalar cases.

[ztangent]

Thanks for making this PR! I'm not sure why the Travis build errored, but separate of that, it'd great if you could adjust the implementation to be overflow-safe in line with #321!

[bzinberg]

Thanks @MrVPlusOne! A couple of initial comments now; I'll come back later to suggest unit / regression tests.

[bzinberg]

Why this change?

[MrVPlusOne]

I'm actually a little unsure how to fix these tests. It seems that the first argument of isapprox is a zero-dimensional array (which I believe is the expected behavior) but the second argument is a Float64, which causes Julia to complain that no method of isapprox matches the given argument types. Any suggestions on how to fix this?

[bzinberg]

After investigating this a bit, I think it boils down to JuliaLang/julia#28866, which I consider to be a bug in Julia. Yes -- in the version of the code before this PR, both the LHS and the RHS should be 0-dimensional arrays. For the RHS, the issue appears to be that these lines

Gen.jl/test/runtests.jl

Lines 22 to 23 in d79aded

	pos_args[i] = copy(args[i]) .+ dx
	neg_args[i] = copy(args[i]) .- dx

do ::Aray{T,0} .+ ::T and ::Aray{T,0} .- ::T, which should give an ::Array{T,0} but due to JuliaLang/julia#28866 gives a scalar. So the LHS of the ./ on the following line should have an Array{T,0} as the ith argument to f, thus (since broadcast = true) should output an Array{T,0}, and consequently the division operation should be ::Array{T,0} ./ ::T which should output an Array{T,0}, but that is not what happens.

I think the cleanest way to fix this is to add a workaround to the (test-only) function finite_diff so that it handles zero-dimensional array valued arguments correctly.

[bzinberg]

(I'm currently drafting this.)

[bzinberg]

Aha, in addition to the above, the function f being supplied to finite_diff does not satisfy the requirements in its docstring: it takes array-valued arguments, but is not the broadcast of a function that takes scalar-valued arguments (i.e. logpdf_grad(vector-valued distribution parameters) is not a broadcast of a bunch of logpdf_grad(scalar-valued distribution parameters)'s). So I guess we should add a variant of finite_diff that has the semantics we need.

[MrVPlusOne]

Yeah, I think to properly test this, we will need a more general finite_diff function or use an autodiff library implementation as a reference.

[bzinberg]

you clearly trust autodiff more than I do 🙂

[bzinberg]

@MrVPlusOne, I've drafted a fix to the unit test breakage. Take a look and LMKWYT?

[MrVPlusOne]

@bzinberg Thanks! The fix looks good to me!

[bzinberg]

@MrVPlusOne, I've made a few more changes that are mostly by way of tidying.

• Rename the vars in the unbroadcast helper functions to generically describe what they do, rather than coupling them to how they're used in logpdf_grad; e.g., unbroadcast_for_arg(arg, grad) becomes _unbroadcast_like(a, full_arr)
• Add some documentation, mainly for what "unbroadcast" means since I can imagine that being ambiguous or confusing to a reader who doesn't have the same context we have
• A bit of code tidying without changing the logic
• Line length and indentation

The one functional change I made was to add to a couple of the unit tests an explicit check on the shape of the returned arrays.

I think regression testing is covered by the shape checks that were fixed in e712c71. They passed before, but that is because the unit tests were incorrect.

@MrVPlusOne, any comments you'd like to make on the above? If looks good to you, I think this is ready to merge.

[MrVPlusOne]

Thanks for the cleaning up; they all look good to me! I think one more nice thing to have would be to add a unit test for the case where the two arguments of broadcasted_norm have different numbers of dimensions. We currently don't have a test for that, right?

[bzinberg]

Great idea @MrVPlusOne, how's this?

[MrVPlusOne]

Thanks, looks good to me!

[bzinberg]

Thanks @MrVPlusOne for using and improving Gen!