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):
-
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 *.
-
This line
|
capture(t::AbstractArray) = istracked(t) ? map!(capture, similar(t), t) : copy(t) |
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:
|
@inline deriv!(t::TrackedArray, v::AbstractArray) = (copyto!(deriv(t), v); nothing) |
I was able to fix these locally by:
-
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.
-
- 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.
- Redefining
deriv!(::StaticArray ::AbstractArray) to apply deriv! elementwise, like it would if the argument were a Tuple.
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):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*.This line
ReverseDiff.jl/src/tracked.jl
Line 225 in 71c5ac0
assumes that
similar(t)supportssetindex!. In the case ofStaticArrays, this assumption breaks ifisbitstype(eltype(t))isfalse, and in particular iftis anStaticArrayofTrackedReals. Similar issue on this line:ReverseDiff.jl/src/tracked.jl
Line 166 in 71c5ac0
I was able to fix these locally by:
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 thesetindex!issue, I had to cast theStaticArrayto anArrayfirst.ReverseDiff.captureonStaticArrays to do the same thing without mutation (usingmap+ array conversion instead ofsimilar+map!). As an aside, if this is equally efficient, seems like we might as well make this the default definition.deriv!(::StaticArray ::AbstractArray)to applyderiv!elementwise, like it would if the argument were aTuple.