I therefore tried to get rid of the Broadcast entirely, hoping that type stability would improve performance. First I tried a custom implementation of pairwise sum for two function arguments (i.e. I implemented sum(f, X, Y), which is equivalent to mapreduce(f, +, X, Y) based on the implementation of mapreduce(f, +, X) in Base, see #77 for the reason behind this). That implementation can be found in here.
Although that implementation makes the function type stable, it was still not very performant. For this reason, in #90 I chose an implementation which allocates an explicit array, which I believed to be more Zygote-friendly. The large improvements over the old versions seen in the benchmarks confirmed this intuition (see https://github.com/JuliaGaussianProcesses/GPLikelihoods.jl/pull/90#issuecomment-1163515226).
However, the solution is not very clean, as it pays for this performance improvement with additional allocations in the forward pass. It is also unclear whether this implementation will be as beneficial or generalize to other AD backends. The potentially better approach would be to define an rrule for the broadcasted sum, as suggested in https://github.com/JuliaGaussianProcesses/GPLikelihoods.jl/pull/90#issuecomment-1163548553 (although even then it would be nice to have the function be type-stable anyway).
For convenience, the old version of
expected_loglikelihood(Gauss-Hermite quadrature method) looked like this: https://github.com/JuliaGaussianProcesses/GPLikelihoods.jl/blob/e9b7da99e46f56859209ff27bbb36d12512d0ad4/src/expectations.jl#L83-L10990 introduced a work-around/hack for two (possibly interrelated) issues of that implementation:
NegativeBinomialLikelihood,The type-instability of the broadcast could be related to https://github.com/JuliaLang/julia/issues/45748 (see also https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/issues/458, which documents a strange behavior in which inference depends on previous evaluations, which I also observed in the
Broadcastconstruction, but could not resolve).I therefore tried to get rid of the
Broadcastentirely, hoping that type stability would improve performance. First I tried a custom implementation of pairwise sum for two function arguments (i.e. I implementedsum(f, X, Y), which is equivalent tomapreduce(f, +, X, Y)based on the implementation ofmapreduce(f, +, X)in Base, see #77 for the reason behind this). That implementation can be found in here.Although that implementation makes the function type stable, it was still not very performant. For this reason, in #90 I chose an implementation which allocates an explicit array, which I believed to be more Zygote-friendly. The large improvements over the old versions seen in the benchmarks confirmed this intuition (see https://github.com/JuliaGaussianProcesses/GPLikelihoods.jl/pull/90#issuecomment-1163515226).
However, the solution is not very clean, as it pays for this performance improvement with additional allocations in the forward pass. It is also unclear whether this implementation will be as beneficial or generalize to other AD backends. The potentially better approach would be to define an
rrulefor the broadcasted sum, as suggested in https://github.com/JuliaGaussianProcesses/GPLikelihoods.jl/pull/90#issuecomment-1163548553 (although even then it would be nice to have the function be type-stable anyway).