Pullback for tr produces a CPU Diagonal causing downstream scalar indexing on GPUs

[darsnack]

Functions like tr(A * B) will throw scalar indexing issues in the pullback for * when A and B are CuArrays. This is because the pullback for tr creates a Diagonal which will cause downstream matrix multiplies to hit the LinearAlgebra definition.

[mcabbott]

Worse it's a Diagonal{T,Array}. Would a Diagonal{T, CuArray} work?

[darsnack]

Yeah, that seems to avoid scalar indexing

[mcabbott]

Then probably it can re-use what sum does, which should also allow 2nd derivatives:

ChainRules.jl/src/rulesets/Base/mapreduce.jl

Line 47 in a9a84ba

_unsum(x, dy, ::Colon) = broadcast(last∘tuple, x, Ref(dy))

It would also be nice if the test noticed this. We have @gpu test_rrule(tr, randn(4, 4)) but it apparently isn't smart enough to object to the Array.