Faster arraydist with LazyArrays.jl

src/DistributionsAD.jl

@@ -80,45 +80,7 @@ include("zygote.jl")
     end
 
     @require LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" begin
-        using .LazyArrays: BroadcastArray, BroadcastVector, LazyArray
-
-        const LazyVectorOfUnivariate{
-            S<:ValueSupport,
-            T<:UnivariateDistribution{S},
-            Tdists<:BroadcastVector{T},
-        } = VectorOfUnivariate{S,T,Tdists}
-
-        function Distributions._logpdf(
-            dist::LazyVectorOfUnivariate,
-            x::AbstractVector{<:Real},
-        )
-            return sum(copy(logpdf.(dist.v, x)))
-        end
-
-        function Distributions.logpdf(
-            dist::LazyVectorOfUnivariate,
-            x::AbstractMatrix{<:Real},
-        )
-            size(x, 1) == length(dist) ||
-                throw(DimensionMismatch("Inconsistent array dimensions."))
-            return vec(sum(copy(logpdf.(dists, x)), dims = 1))
-        end
-
-        const LazyMatrixOfUnivariate{
-            S<:ValueSupport,
-            T<:UnivariateDistribution{S},
-            Tdists<:BroadcastArray{T,2},
-        } = MatrixOfUnivariate{S,T,Tdists}
-
-        function Distributions._logpdf(
-            dist::LazyMatrixOfUnivariate,
-            x::AbstractMatrix{<:Real},
-        )
-            return sum(copy(logpdf.(dist.dists, x)))
-        end
-
-        lazyarray(f, x...) = LazyArray(Base.broadcasted(f, x...))
-        export lazyarray
+        include("lazyarrays.jl")
     end
 end
 

src/common.jl

@@ -48,3 +48,112 @@ parameterless_type(x) = parameterless_type(typeof(x))
 parameterless_type(x::Type) = __parameterless_type(x)
 
 @non_differentiable adapt_randn(::Any...)
+
+"""
+    make_closure(f, g)
+
+Return a closure of the form `(x, args...) -> f(g(args...), x)`.
+
+# Examples
+
+This is particularly useful when one wants to avoid broadcasting over constructors
+which can sometimes cause issues with type-inference, in particular when combined
+with reverse-mode AD frameworks.
+
+```juliarepl
+julia> using DistributionsAD, Distributions, ReverseDiff, BenchmarkTools
+
+julia> const data = randn(1000);
+
+julia> x = randn(length(data));
+
+julia> f(x) = sum(logpdf.(Normal.(x), data))
+f (generic function with 2 methods)
+
+julia> @btime ReverseDiff.gradient(\$f, \$x);
+  848.759 μs (14605 allocations: 521.84 KiB)
+
+julia> # Much faster with ReverseDiff.jl.
+       g(x) = let g_inner = DistributionsAD.make_closure(logpdf, Normal)
+           sum(g_inner.(data, x))
+       end
+g (generic function with 1 method)
+
+julia> @btime ReverseDiff.gradient(\$g, \$x);
+  17.460 μs (17 allocations: 71.52 KiB)
+```
+
+See https://github.com/TuringLang/Turing.jl/issues/1934 more further discussion.
+
+# Notes
+To really go "vrooom!\" one needs to specialize on the arguments, e.g. if one
+has a function `myfunc` then we need to define
+
+```julia
+make_closure(::typeof(myfunc), ::Type{D}) where {D} = myfunc(D(args...), x)
+```
+
+This can also be done using `DistributionsAD.@specialize_make_closure`:
+
+```julia
+julia> mylogpdf(d, x) = logpdf(d, x)
+mylogpdf (generic function with 1 method)
+
+julia> h(x) = let inner = DistributionsAD.make_closure(mylogpdf, Normal)
+           sum(inner.(data, x))
+       end
+h (generic function with 1 method)
+
+julia> @btime ReverseDiff.gradient(\$h, \$x);
+  1.220 ms (37011 allocations: 1.42 MiB)
+
+julia> DistributionsAD.@specialize_make_closure mylogpdf
+
+julia> @btime ReverseDiff.gradient(\$h, \$x);
+  17.038 μs (17 allocations: 71.52 KiB)
+```
+"""
+make_closure(f, g) = (x, args...) -> f(g(args...), x)
+make_closure(f, ::Type{D}) where {D} = (x, args...) -> f(D(args...), x)
+
+
+"""
+    has_specialized_make_closure(f, g)
+
+Return `true` if there exists a specialized `make_closure(f, g)` implementation.
+"""
+has_specialized_make_closure(f, g) = false
+
+# To go vroooom we need to specialize on the first argument, thus ensuring that
+# a different closure is constructed for each method.
+"""
+    @specialize_make_closure(f)
+
+Define `make_closure` and `has_specialized_make_closure` for first first argument being `f` 
+and second argument being a type.
+"""
+macro specialize_make_closure(f)
+    return quote
+        $(DistributionsAD).make_closure(::typeof($(esc(f))), ::Type{D}) where {D} = (x, args...) -> $(esc(f))(D(args...), x)
+        $(DistributionsAD).has_specialized_make_closure(::typeof($(esc(f))), ::Type{D}) where {D} = true
+    end
+end
+
+"""
+    @specialize_make_closure(f, g)
+
+Define `make_closure` and `has_specialized_make_closure` for first first argument being `f` 
+and second argument being `g`.
+"""
+macro specialize_make_closure(f, g)
+    return quote
+        $(DistributionsAD).make_closure(::typeof($(esc(f))), ::typeof($(esc(g)))) = (x, args...) -> $(esc(f))($(esc(g))(args...), x)
+        $(DistributionsAD).has_specialized_make_closure(::typeof($(esc(f))), ::typeof{$(esc(g))}) = true
+    end
+end
+
+@specialize_make_closure Distributions.pdf
+@specialize_make_closure Distributions.logpdf
+@specialize_make_closure Distributions.loglikelihood
+@specialize_make_closure Distributions.cdf
+@specialize_make_closure Distributions.logcdf

src/lazyarrays.jl

@@ -0,0 +1,53 @@
+using .LazyArrays: BroadcastArray, BroadcastVector, LazyArray
+
+const LazyVectorOfUnivariate{
+    S<:ValueSupport,
+    T<:UnivariateDistribution{S},
+    Tdists<:BroadcastVector{T},
+} = VectorOfUnivariate{S,T,Tdists}
+
+_inner_constructor(::Type{<:BroadcastVector{<:Any,Type{D}}}) where {D} = D
+
+function Distributions._logpdf(
+    dist::LazyVectorOfUnivariate,
+    x::AbstractVector{<:Real},
+)
+    # TODO: Implement chain rule for `LazyArray` constructor to support Zygote.
+    f = make_closure(logpdf, _inner_constructor(typeof(dist.v)))
+    # TODO: Make use of `sum(Broadcast.instantiate(Broadcast.broadcasted(f, x, args...)))` once
+    # we've addressed performance issues in ReverseDiff.jl.
+    return sum(f.(x, dist.v.args...))
+end
+
+function Distributions.logpdf(
+    dist::LazyVectorOfUnivariate,
+    x::AbstractMatrix{<:Real},
+    )
+    size(x, 1) == length(dist) ||
+        throw(DimensionMismatch("Inconsistent array dimensions."))
+    f = make_closure(logpdf, _inner_constructor(typeof(dist.v)))
+    return vec(sum(f.(x, dist.v.args...), dims = 1))
+end
+
+const LazyMatrixOfUnivariate{
+    S<:ValueSupport,
+    T<:UnivariateDistribution{S},
+    Tdists<:BroadcastArray{T,2},
+} = MatrixOfUnivariate{S,T,Tdists}
+
+function Distributions._logpdf(
+    dist::LazyMatrixOfUnivariate,
+    x::AbstractMatrix{<:Real},
+)
+    f = make_closure(logpdf, _inner_constructor(typeof(dist.v)))
+
+    return sum(f.(x, dist.v.args))
+end
+
+lazyarray(f, x...) = BroadcastArray(f, x...)
+export lazyarray
+
+# Necessary to make `BroadcastArray` work nicely with Zygote.
+function ChainRulesCore.rrule(config::RuleConfig{>:HasReverseMode}, ::Type{LazyArrays.BroadcastArray}, f, args...)
+    return ChainRulesCore.rrule_via_ad(config, Broadcast.broadcasted, f, args...)
+end