JuliaDiff / JuliaDiff/ReverseDiff.jl

ChainRulesCore projection & ReverseDiff.TrackedArray

Open
#204 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
393
Forks
60
Avg merge
18h 24m
Merged PRs (30d)
8

Description

When using ReverseDiff over Zygote, its tracked arrays will often be turned into arrays of TrackedReal:

julia> using ReverseDiff, Zygote

julia> _, back = pullback(x -> cumsum(x .^ 3), rand(3))
([0.0417770851525806, 0.08898338554941161, 0.13448629430576223], Zygote.var"#52#53"{typeof(∂(#11))}(∂(#11)))

julia> ta = ReverseDiff.track([1,2,3.0]);

julia> back(ta)[1]
3-element Vector{ReverseDiff.TrackedReal{Float64, Float64, Nothing}}:
 TrackedReal<Ha5>(2.1671946936904316, 0.0, BcR, ---)
 TrackedReal<J2r>(1.9592562290038964, 0.0, BcR, ---)
 TrackedReal<5TB>(1.147101768209455, 0.0, BcR, ---)

The cause of this is that ChainRulesCore's projection mechanism looks at the eltype, and decides that it needs correcting:

julia> using ChainRulesCore

julia> pt = ProjectTo(ta)
ProjectTo{AbstractArray}(element = ProjectTo{Real}(), axes = (Base.OneTo(3),))

julia> pt(ta)  # this works fine, TrackedArray
3-element ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}:
 TrackedReal<EIm>(1.0, 0.0, BcR, 1, 96O)
 TrackedReal<I8K>(2.0, 0.0, BcR, 2, 96O)
 TrackedReal<9lo>(3.0, 0.0, BcR, 3, 96O)

julia> pr = ProjectTo(rand(3))
ProjectTo{AbstractArray}(element = ProjectTo{Float64}(), axes = (Base.OneTo(3),))

julia> pr(ta)  # makes a Vector{TrackedReal}
3-element Vector{ReverseDiff.TrackedReal{Float64, Float64, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}}}:
 TrackedReal<KBA>(1.0, 0.0, BcR, 1, 96O)
 TrackedReal<29L>(2.0, 0.0, BcR, 2, 96O)
 TrackedReal<7oX>(3.0, 0.0, BcR, 3, 96O)

julia> map(pr.element, ta)  # this is what projection calls, as TrackedReal <: Float64 is false
3-element Vector{ReverseDiff.TrackedReal{Float64, Float64, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}}}:
 TrackedReal<AQQ>(1.0, 0.0, BcR, 1, 96O)
 TrackedReal<5IS>(2.0, 0.0, BcR, 2, 96O)
 TrackedReal<2yC>(3.0, 0.0, BcR, 3, 96O)

I think this should probably be avoided, by adding methods to (::ProjectTo)(::TrackedArray) which bypass this behaviour. Since ReverseDiff already depends on CRC, no extra deps are required.

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 reading ChainRulesCore's ProjectTo dispatch and ReverseDiff.TrackedArray definitions, then reproduce the shown ProjectTo(ta) and ProjectTo(rand(3))(ta) behavior. Done means projecting a TrackedArray preserves its tracked-array form instead of producing an array of TrackedReal, while ordinary array projection continues to work.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.