Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Accumulation #69

Open
mcabbott opened this issue Jan 11, 2022 · 0 comments
Open

Accumulation #69

mcabbott opened this issue Jan 11, 2022 · 0 comments

Comments

@mcabbott
Copy link
Member

mcabbott commented Jan 11, 2022

At present Diffractor may return thunks, but doesn't seem use them (or anything else) to accumulate effciently:

julia> @btime gradient(x -> sum(x), $(rand(100, 100))) |> first |> typeof
  min 1.829 μs, mean 1.869 μs (8 allocations, 368 bytes)
ChainRulesCore.InplaceableThunk{ChainRulesCore.Thunk...

julia> @btime gradient(x -> sum(x) + sum(x), $(rand(100, 100))) |> first |> typeof
  min 8.917 μs, mean 28.032 μs (22 allocations, 235.03 KiB)
Matrix{Float64} (alias for Array{Float64, 2})

julia> @btime copy($(rand(100, 100)));
  min 1.262 μs, mean 4.120 μs (2 allocations, 78.17 KiB)

julia> 235.03 / 78.17
3.00665216835103

Should this change, perhaps to use ChainRulesCore.add!!? In which case it might it be easiest to change now, while there is nothing downstream to break.

Note aside that add!! is slower than expected here, 3 copies not 1:

julia> g1 = gradient(x -> sum(x), (rand(100, 100)))[1]; g2 = deepcopy(g1);

julia> @btime ChainRulesCore.add!!($(g1), $(g2));
  min 4.719 μs, mean 21.300 μs (8 allocations, 234.55 KiB)

julia> 234.55 / 78.17
3.0005117052577717

julia> @btime ChainRulesCore.add!!($(randn(100, 100)), $(g2));
  min 3.318 μs, mean 3.438 μs (0 allocations)

Or do we count on ImmutableArray and compiler improvements?

Xref FluxML/Zygote.jl#981 which retrofits Zygote to accumulate in-place.

Xref also JuliaDiff/ChainRulesCore.jl#539 which alters + on two thunks to be more efficient.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant