Operations on StaticArrays give runtime BoundsError

[platawiec]

Possibly related to #837

This is on Julia 1.4, Zygote 0.6.0, and StaticArrays 0.12.5

using Zygote
using StaticArrays
using LinearAlgebra

Zygote.gradient(x -> norm(x), SVector(1.0, 1.0))

Beginning of stacktrace:

Internal error: encountered unexpected error in runtime:
BoundsError(a=Array{Any, (2,)}[
  Core.Compiler.VarState(typ=Zygote.Pullback{Tuple{typeof(StaticArrays._norm), StaticArrays.Size{(2,)}, StaticArrays.SArray{Tuple{2}, Float64, 1, 2}}, Any}, undef=false),
  Core.Compiler.VarState(typ=Float64, undef=false)], i=(3,))
jl_bounds_error_ints at /Users/julia/buildbot/worker/package_macos64/build/src/rtutils.c:183
stupdate! at /Applications/Julia-1.4.app/Contents/Resources/julia/lib/julia/sys.dylib (unknown line)
...

[DhairyaLGandhi]

Likely related to chain rules rules being too conservative

[platawiec]

I did a bit of sleuthing, the above works for Zygote v0.5.15 , but fails for v0.5.16 . Specifically, #847 likely triggered this and I ran into it because I needed norm.

[DhairyaLGandhi]

Can you revert that pr and test? If it works we can figure out supporting svectors in Zygote directly like we used to.

[platawiec]

According to the changelog, that PR was the only change between 0.5.15 and 0.5.16.

Another approach would be widening the types allowed in ChainRulesCore (to e.g. AbstractVector), but I'm sure there's a discussion I missed that goes over why that's a bad idea.

[mcabbott]

The ChainRules pullback looks like it would work fine, I'm not sure why it was written for StridedArray. (There is a long debate somewhere... but this case seems pretty clear.)

[oxinabox]

Zygote should decompose the operation on the StaticArray and find on the inside something that it can work with.
we can relax the rule in ChainRules as a hack for now though.