JuliaDiff / JuliaDiff/ReverseDiff.jl

`StaticArray` of `TrackedReal`s doesn't currently work (and I think it could)

Open
#153 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

Sorry for the lack of minimal example here, I'll add it when I get the chance.

I recently tried to differentiate through operations on StaticArrays.SVectors and hit two issues, both of which were easily fixable once diagnosed (unless these solutions cause other problems I'm not aware of):

  1. StaticArrays defines a catch-all +(::StaticArray, ::AbstractArray) and ReverseDiff defines a catch-all +(::AbstractArray, ::TrackedArray). So if +(::StaticArray, ::TrackedArray) is called, Julia complains that the dispatch rule is ambiguous. (Same for the other order of operands too.) The same is the case for -, and I'm guessing also for *.

  2. This line
    https://github.com/JuliaDiff/ReverseDiff.jl/blob/71c5ac0cb2212f49e3b3d567b45a830614173913/src/tracked.jl#L225
    assumes that similar(t) supports setindex!. In the case of StaticArrays, this assumption breaks if isbitstype(eltype(t)) is false, and in particular if t is an StaticArray of TrackedReals. Similar issue on this line:
    https://github.com/JuliaDiff/ReverseDiff.jl/blob/71c5ac0cb2212f49e3b3d567b45a830614173913/src/tracked.jl#L166

I was able to fix these locally by:

  1. Explicitly defining +(::StaticArray, ::TrackedArray) and +(::TrackedArray, ::StaticArray), to disambiguate the dispatch rulea similarly for - (and probably also for *, though I didn't have the need for that). In order to not encounter the setindex! issue, I had to cast the StaticArray to an Array first.

    1. Redefining ReverseDiff.capture on StaticArrays to do the same thing without mutation (using map + array conversion instead of similar + map!). As an aside, if this is equally efficient, seems like we might as well make this the default definition.
    2. Redefining deriv!(::StaticArray ::AbstractArray) to apply deriv! elementwise, like it would if the argument were a Tuple.

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 at src/tracked.jl lines 166 and 225, then reproduce the reported failures with StaticArrays.SVector and StaticArray values containing TrackedReal elements. Check the relevant StaticArrays and ReverseDiff method interactions; done means supported addition and subtraction dispatches no longer conflict and capture or deriv! handles these arrays without assuming mutable similar storage.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.