Faster filldist() #227

alyst · 2022-07-06T23:53:42Z

Switch logpdf(::FillDist) from sum(map(f, x)) to mapreduce(f, +, x) to eliminate the unnecessary array allocations.

sum(f, x) would have been even simpler, but I'm hitting *(::TrackedReal, ::Dual) method ambiguity on this path.

It's beyond my current understanding of the autodiff code to figure out whether just this change breaks some tests (Chernoff distribution) and fixes other (Zygote) or it is just a coincidence.

I have tried to @benchmark the change, but with @benchmark(logpdf($(filldist(Cauchy(), 1000)), $(rand(1000)))) I don't see much of the change.
I guess it starts to be visible with dual numbers.

So here's the profiling results.

Before:

After:

Note that array allocation and copying disappear, and it looks like overall it should be ~30% faster.

define vars used by error()

devmotion

Thanks for the PR. I guess these methods are mainly optimized for compatibility with most AD backends (and in particular Zygote) and hence for historical reasons used whatever was possible and inferred at the time when they were introduced.

(A side remark: The long term plan is to remove all of this functionality from DistributionsAD and generalize it in Distributions: JuliaStats/Distributions.jl#1391)

src/arraydist.jl

src/filldist.jl

Co-authored-by: David Widmann <[email protected]>

alyst · 2022-07-20T21:15:02Z

Is there any chance to merge it?

devmotion · 2022-07-20T21:29:59Z

Sure, if tests pass successfully I assume it can be approved by a maintainer and merged eventually.

alyst · 2022-07-20T21:45:10Z

Sure, if tests pass successfully I assume it can be approved by a maintainer and merged eventually.

I hope the fact that the tests got cancelled for some backends is some temporary issue that is not related to this PR. Otherwise it would be hard to know whether the tests pass.

devmotion · 2022-07-20T21:51:37Z

src/filldist.jl

@@ -30,21 +30,18 @@ end
 function _flat_logpdf(dist, x)
    if toflatten(dist)
        f, args = flatten(dist)
-        return sum(f.(args..., x))
+        return mapreduce(xi -> f(args..., xi), +, x)


A bit simpler:

Suggested change

return mapreduce(xi -> f(args..., xi), +, x)

return sum(xi -> f(args..., xi), x)

Let's try, I hope it would not break Tracker.

Alas, this triggers Tracker method ambiguity

Multivariate distributions: Error During Test at C:\Users\astukalov\.julia\dev\DistributionsAD\test\ad\utils.jl:357 Test threw exception Expression: ≈(Tracker.data((Tracker.gradient(f, x))[1]), finitediff, rtol = rtol, atol = atol) MethodError: -(::ForwardDiff.Dual{Nothing, Float64, 1}, ::Tracker.TrackedReal{Float64}) is ambiguous. Candidates: -(a::Real, b::Tracker.TrackedReal) in Tracker at C:\Users\astukalov\.julia\packages\Tracker\9xWLl\src\lib\real.jl:96 -(x::ForwardDiff.Dual{Tx}, y::Real) where Tx in ForwardDiff at C:\Users\astukalov\.julia\packages\ForwardDiff\wAaVJ\src\dual.jl:144 Possible fix, define -(::ForwardDiff.Dual{Tx}, ::Tracker.TrackedReal) where Tx Stacktrace: [1] zval(μ::Tracker.TrackedReal{Float64}, σ::Tracker.TrackedReal{Float64}, x::ForwardDiff.Dual{Nothing, Float64, 1}) @ StatsFuns C:\Users\astukalov\.julia\packages\StatsFuns\vxSkw\src\distrs\norm.jl:10 [2] normlogpdf(μ::Tracker.TrackedReal{Float64}, σ::Tracker.TrackedReal{Float64}, x::ForwardDiff.Dual{Nothing, Float64, 1}) @ StatsFuns C:\Users\astukalov\.julia\packages\StatsFuns\vxSkw\src\distrs\norm.jl:39 [3] logpdf @ C:\Users\astukalov\.julia\packages\Distributions\QLJcf\src\univariate\continuous\logitnormal.jl:126 [inlined] [4] #200 @ .\none:0 [inlined] [5] #21 @ C:\Users\astukalov\.julia\dev\DistributionsAD\src\filldist.jl:33 [inlined] [6] partial(f::DistributionsAD.var"#21#22"{Tuple{Tracker.TrackedReal{Float64}, Tracker.TrackedReal{Float64}}, DistributionsAD.var"#200#201"}, Δ::Tracker.TrackedReal{Float64}, i::Int64, args::Float64) @ Tracker C:\Users\astukalov\.julia\packages\Tracker\9xWLl\src\lib\array.jl:573 [7] _broadcast_getindex_evalf @ .\broadcast.jl:670 [inlined] [8] _broadcast_getindex @ .\broadcast.jl:643 [inlined] [9] getindex @ .\broadcast.jl:597 [inlined] [10] copy @ .\broadcast.jl:899 [inlined] [11] materialize(bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(Tracker.partial), Tuple{Base.RefValue{DistributionsAD.var"#21#22"{Tuple{Tracker.TrackedReal{Float64}, Tracker.TrackedReal{Float64}}, DistributionsAD.var"#200#201"}}, Vector{Tracker.TrackedReal{Float64}}, Int64, Vector{Float64}}}) @ Base.Broadcast .\broadcast.jl:860 [12] broadcast(::typeof(Tracker.partial), ::Base.RefValue{DistributionsAD.var"#21#22"{Tuple{Tracker.TrackedReal{Float64}, Tracker.TrackedReal{Float64}}, DistributionsAD.var"#200#201"}}, ::Vector{Tracker.TrackedReal{Float64}}, ::Int64, ::Vararg{Any}) @ Base.Broadcast .\broadcast.jl:798 [13] ∇broadcast @ C:\Users\astukalov\.julia\packages\Tracker\9xWLl\src\lib\array.jl:577 [inlined] [14] copy(bc::Base.Broadcast.Broadcasted{Tracker.TrackedStyle, Tuple{Base.OneTo{Int64}}, typeof(Tracker.partial), Tuple{Base.RefValue{DistributionsAD.var"#21#22"{Tuple{Tracker.TrackedReal{Float64}, Tracker.TrackedReal{Float64}}, DistributionsAD.var"#200#201"}}, Vector{Tracker.TrackedReal{Float64}}, Int64, TrackedArray{…,Vector{Float64}}}}) @ Tracker C:\Users\astukalov\.julia\packages\Tracker\9xWLl\src\lib\array.jl:607 [15] materialize @ .\broadcast.jl:860 [inlined] [16] (::Tracker.var"#621#623"{Vector{Tracker.TrackedReal{Float64}}, DistributionsAD.var"#21#22"{Tuple{Tracker.TrackedReal{Float64}, Tracker.TrackedReal{Float64}}, DistributionsAD.var"#200#201"}, Tuple{TrackedArray{…,Vector{Float64}}}})(i::Int64) @ Tracker C:\Users\astukalov\.julia\packages\Tracker\9xWLl\src\lib\array.jl:580

I guess it has to be fixed on the Tracker side. Also looks like it still generates an array.

I guess it has to be fixed on the Tracker side. Also looks like it still generates an array.

Well, Tracker defines sum(f, xs::TrackedArray) as sum(f.(xs)): https://github.com/FluxML/Tracker.jl/blob/84ff74daaa250dadb424d251c66d64fa64ade819/src/lib/array.jl#L358 So it will generate an intermediate array and will use ForwardDiff when differentiating (as broadcasting operations with Tracker involve ForwardDiff).

src/arraydist.jl

devmotion · 2022-07-20T21:59:32Z

src/filldist.jl

-        return sum(map(x) do x
-            logpdf(dist, x)
-        end)
+        return mapreduce(Base.Fix1(logpdf, dist), +, x)


Suggested change

return mapreduce(Base.Fix1(logpdf, dist), +, x)

return sum(Base.Fix1(logpdf, dist), x)

devmotion · 2022-07-20T22:00:09Z

src/filldist.jl

    end
 end

 function _flat_logpdf_mat(dist, x)
    if toflatten(dist)
        f, args = flatten(dist)
-        return vec(sum(f.(args..., x), dims = 1))
+        return vec(mapreduce(xi -> f(args..., xi), +, x, dims = 1))


Suggested change

return vec(mapreduce(xi -> f(args..., xi), +, x, dims = 1))

return vec(sum(xi -> f(args..., xi), x; dims=1))

devmotion · 2022-07-20T22:00:24Z

src/filldist.jl

    else
-        temp = map(x -> logpdf(dist, x), x)
-        return vec(sum(temp, dims = 1))
+        return vec(mapreduce(Base.Fix1(logpdf, dist), +, x; dims = 1))


Suggested change

return vec(mapreduce(Base.Fix1(logpdf, dist), +, x; dims = 1))

return vec(sum(Base.Fix1(logpdf, dist), x; dims = 1))

test/ad/distributions.jl

Co-authored-by: David Widmann <[email protected]>

since it looks like `mapreduce()` still allocates Co-authored-by: David Widmann <[email protected]>

some graident tests require test_approx(::Array{<:Array}, ::Zero)

alyst · 2022-08-04T21:41:12Z

This is my best attempt to make the tests pass. They take quite long (that is why CI times out), and actually don't complete. Here's the message I'm getting:

[ Info: Testing: arraydist(MvNormal, 2)
[ Info: Testing: filldist(MvNormal, 2, 2)
[ Info: Testing: arraydist(MvNormal, 2, 2)
ERROR: Package DistributionsAD errored during testing (exit code: 3221225727)

I'm not sure whether it is related to the PR though.
Otherwise there tests pass except for Chernoff distribution, which I don't know how to fix, because it looks like some tests for Chernoff distribution pass and some don't.
(And, frankly, the current design of enabling/disabling specific tests is not very straightforward for the newcomer to digest)

devmotion · 2022-08-08T11:18:32Z

src/arraydist.jl

@@ -3,7 +3,7 @@
 const VectorOfUnivariate = Distributions.Product

 function arraydist(dists::AbstractVector{<:UnivariateDistribution})
-    return Product(dists)
+    return product_distribution(dists)


This changes the return type and breaks tests, as observed in #228 (comment).

Let's see whether #228 fixes broken tests.

src/arraydist.jl

test/ad/distributions.jl

yebai · 2022-08-08T19:01:02Z

Many thanks, @alyst and @devmotion!

alyst added 4 commits July 6, 2022 11:45

fix testset_zygote_broken()

345bbcd

define vars used by error()

logpdf(arraydist): use mapreduce

4fd90ea

logpdf(filldist): use mapreduce

6aa1959

remove filldist(Zygote) from broken

d066986

devmotion reviewed Jul 7, 2022

View reviewed changes

src/arraydist.jl Outdated Show resolved Hide resolved

src/filldist.jl Outdated Show resolved Hide resolved

src/filldist.jl Outdated Show resolved Hide resolved

alyst and others added 2 commits July 6, 2022 17:53

improve mapreduce

532f155

Co-authored-by: David Widmann <[email protected]>

improve mapreduce() invocation

b4c47af

Co-authored-by: David Widmann <[email protected]>

tests: exclude Chernoff from Zygote filldist tests

ed1471e

devmotion reviewed Jul 20, 2022

View reviewed changes

alyst and others added 2 commits July 20, 2022 15:31

simplify mapreduce -> sum

c914d27

Co-authored-by: David Widmann <[email protected]>

explicitly broadcast

867651e

since it looks like `mapreduce()` still allocates Co-authored-by: David Widmann <[email protected]>

alyst mentioned this pull request Jul 26, 2022

Add ZeroTangent vs array of arrays tests JuliaDiff/ChainRulesTestUtils.jl#257

Merged

alyst added 7 commits July 30, 2022 12:55

require ChainRulesTestUtils >= 1.9.2

6690440

some graident tests require test_approx(::Array{<:Array}, ::Zero)

_flat_logpdf(): explicit lazy broadcasting

bea4b35

filldist tests: enable Skellam

bc6c50a

use product_distribution() to fix deprecation

e9226a9

eliminate unnecessary intermediate var

2320bb6

replace some anonymous funcs with Base.Fix1

336aa8b

replace sum(lambda, zip(...)) with lazy broadcast

78047dc

devmotion reviewed Aug 8, 2022

View reviewed changes

yebai reviewed Aug 8, 2022

View reviewed changes

src/arraydist.jl Outdated Show resolved Hide resolved

Update src/arraydist.jl

634ce91

devmotion reviewed Aug 8, 2022

View reviewed changes

test/ad/distributions.jl Outdated Show resolved Hide resolved

Update test/ad/distributions.jl

f2c4bca

yebai merged commit 02ca329 into TuringLang:master Aug 8, 2022

	return mapreduce(xi -> f(args..., xi), +, x)
	return sum(xi -> f(args..., xi), x)

	return mapreduce(Base.Fix1(logpdf, dist), +, x)
	return sum(Base.Fix1(logpdf, dist), x)

	return vec(mapreduce(xi -> f(args..., xi), +, x, dims = 1))
	return vec(sum(xi -> f(args..., xi), x; dims=1))

	return vec(mapreduce(Base.Fix1(logpdf, dist), +, x; dims = 1))
	return vec(sum(Base.Fix1(logpdf, dist), x; dims = 1))

Faster filldist() #227

Faster filldist() #227

Uh oh!

Conversation

alyst commented Jul 6, 2022

Uh oh!

devmotion left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alyst commented Jul 20, 2022

Uh oh!

devmotion commented Jul 20, 2022

Uh oh!

alyst commented Jul 20, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

alyst commented Aug 4, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

yebai commented Aug 8, 2022

Uh oh!

Uh oh!