JuliaDiff / JuliaDiff/ChainRulesCore.jl

summing many differentials (varadic +)

Open
#275 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement inplace accumulation
Dominant language
Julia
Stars
267
Forks
66
PR merge metrics
No merged PRs in 30d

Description

Zygote regularly emits things like

accum(a, b, c, d)

which is equivalent to

+(a, b, c, d)

in ChainRules language. This means that we a) need to support this and b) could optimise it 🥳 .

A simple default implementation already exists provided that +(a, b) is defined, and for simple things like Arrays this is optimised already. For example:

using BenchmarkTools
a = randn(10);
b = randn(10);
c = randn(10);
@benchmark $a + $b + $c

yields

BenchmarkTools.Trial:
  memory estimate:  160 bytes
  allocs estimate:  1
  --------------
  minimum time:     51.717 ns (0.00% GC)
  median time:      52.820 ns (0.00% GC)
  mean time:        55.692 ns (1.67% GC)
  maximum time:     443.181 ns (82.46% GC)
  --------------
  samples:          10000
  evals/sample:     980

Only a single temporary was allocated, which is great. However, we don't have this optimisation for e.g. Composites. Consider

julia> using ChainRulesCore

julia> a = Composite{Any}(a);

julia> b = Composite{Any}(b);

julia> c = Composite{Any}(c);

julia> @benchmark $a + $b + $c
BenchmarkTools.Trial:
  memory estimate:  320 bytes
  allocs estimate:  2
  --------------
  minimum time:     84.682 ns (0.00% GC)
  median time:      88.059 ns (0.00% GC)
  mean time:        93.438 ns (2.04% GC)
  maximum time:     644.918 ns (85.62% GC)
  --------------
  samples:          10000
  evals/sample:     947

It would be nice if we could obtain the same benefits here as in the Array case.

Loosely related to #226 and #113

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the existing default implementation for binary +(a, b) and how Composite addition is handled. Reproduce the provided BenchmarkTools comparisons for Arrays and Composites, then investigate how variadic + could preserve the requested ChainRules behavior while reducing Composite allocations. Done means variadic addition is supported and Composite benchmarks approach the single-temporary behavior described.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.