In general, Zygote tends to “lose” StaticArray gradients because none of its internals are aware of them. That means something as pedestrian as gradient(x -> x[1] + sum(x), SA[1,2,3]) could (and does) end up on an allocating path and returning a SizedArray. Have you considered alternate AD libraries like ReverseDiff?
1 Like