The error might be avoided by sum(abs2, ...) instead of sum.
Zygote’s rule for sum uses FillArrays, which this PR may remove. It’s an optimisation to save a little allocation which causes all kinds of issues (normally in toy problems like this). Here @tensor doesn’t understand this type of fake array, see e.g. this question.