`getindex` thunk: `add!` disagrees with `val`; `sum(f, xs)` accumulates closure tangents quadratically

Open
  • Is something wrong?
#863 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
25/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Domain
machine-learning

Research direction

Start with the thunked_∇getindex reverse rule for getindex, comparing what its add! does with unthunk(dy) against what val does through ProjectTo(x). Then look at the general branch of rrule for sum(f, xs) and its sum(first, ...) accumulation, which needs an add!! that recurses into Tangent fields. Done means the reproducers in the issue give the same result from add! as from unthunk, and the quadratic allocation case runs in linear time. The design choices for the fix need maintainer input first.

Written by the indexing model from the issue text.

Description

Two things in the reverse rules keep an AD from accumulating getindex cotangents in place. That makes x[i] in a loop quadratic (FluxML/Zygote.jl#365). Zygote now accumulates in place (FluxML/Zygote.jl#1738), but these two cases stay quadratic because of the rules here. Both are unchanged on master (v1.74.8).

1. thunked_∇getindex: add! adds raw dy, val projects it

val is ∇getindex(x, dy, inds...), which ends in ProjectTo(x)(dx), but add! adds unthunk(dy) as is. So t.add!(dx) can differ from dx + unthunk(t), or throw:

using ChainRules, ChainRulesCore
x = [1.0, 2.0, 3.0]
dx = rrule(getindex, x, 1)[2](1.0 + 2.0im)[2]
unthunk(dx)                    # [1.0, 0.0, 0.0]
dx.add!(zeros(3))              # InexactError
dx.add!(zeros(ComplexF64, 3))  # [1.0 + 2.0im, 0, 0]

dx = rrule(getindex, x, 1:2)[2]([1.0im, 2.0im])[2]
dx.add!(zeros(3))              # InexactError, though unthunk(dx) == zeros(3)

x32 = Float32[1, 2, 3]
dx = rrule(getindex, x32, 2)[2](0.1)[2]
dx.add!(zeros(3))[2] == unthunk(dx)[2]  # false: 0.1 vs Float32(0.1)

So an AD can't call add! without knowing dy, which is hidden in the closure, and has to fall back to unthunk. For a slice that builds a dense array the size of the whole x each time. A loop over the columns of a 100×1000 matrix then allocates 767 MiB in Zygote.

Suggested fix: make add! add the same projected value as val, at least for plain Array{<:Number}, where ProjectTo(x) is elementwise. For example, dx -> ∇getindex!(dx, _project_dy(x, unthunk(dy)), Base.to_indices(x, inds)...), where _project_dy applies ProjectTo(zero(eltype(x))) to dy, elementwise for slices. Structured x could keep dx .+= unthunk(val). With that patch, Zygote calling add! when the inferred eltype of val fits its buffer takes that loop from 767 MiB to 4.4 MiB.

A related cost: in the last example unthunk(dx) is a dense Vector{Float32}, not a OneElement. Projecting the OneElement onto a different eltype materializes it, so mixed-precision scalar indexing is O(length(x)) per access. Projecting dy before building the OneElement would avoid that.

2. sum(f, xs) adds closure tangents with sum(first, df_and_dxs)

In the general branch of rrule(::RuleConfig{>:HasReverseMode}, ::typeof(sum), f, xs), each + of two Tangents whose fields are InplaceableThunks unthunks and copies the whole array. For a closure that captures x and reads x[i], that is O(n²):

using Zygote
x = rand(10_000)
@allocated gradient(x -> sum(i -> x[i]^2, eachindex(x)), x)  # ≈ 766 MiB

Generators hit the same path, since sum(x[i]^2 for i in ...) is lowered to sum(f, iter). Accumulating df into a partial sum the rule owns, instead of sum(first, ...), would make it linear. That needs an add!! that recurses into Tangent fields, since add!! falls back to + for Tangent. Zygote can route around this through its own map adjoint (766 MiB → 3.8 MiB with #1738), but other ADs using this rule would still pay.

🧔‍♂️🤖 Assisted by AI

Dominant language
Julia
Stars
479
Forks
100
Avg merge
1d 4h
Merged PRs (30d)
14

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

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.

More from JuliaDiff/ChainRules.jl

All issues in JuliaDiff/ChainRules.jl

Similar issues

More Julia issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.