Grads algebra on GPU causes scalar getindex is disallowed error

[pxl-th]

Hi! I was trying to do some gradient accumulation, but Grads algebra seems to cause error on GPU.
Am I missing something or is there indeen an issue?

I was using it in the context of UNet model training, but below is MWE.:

using CUDA
CUDA.allowscalar(false)
using Zygote

w = randn(Float32, 2) |> cu
x1 = randn(Float32, 2) |> cu
x2 = randn(Float32, 2) |> cu
p = Params([w])

g1 = gradient(() -> sum(w .* x1), p)
g2 = gradient(() -> sum(w .* x2), p)

g1 .+= g2

Error:

ERROR: scalar getindex is disallowed
Stacktrace:
  [1] error(s::String)
    @ Base .\error.jl:33
  [2] assertscalar(op::String)
    @ GPUArrays C:\Users\tonys\.julia\packages\GPUArrays\4n0iS\src\host\indexing.jl:62
  [3] getindex
    @ C:\Users\tonys\.julia\packages\GPUArrays\4n0iS\src\host\indexing.jl:104 [inlined]
  [4] hash(A::CuArray{Float32, 1}, h::UInt64)
    @ Base .\abstractarray.jl:2432
  [5] hash
    @ .\hashing.jl:18 [inlined]
  [6] hashindex
    @ .\dict.jl:169 [inlined]
  [7] ht_keyindex2!(h::Dict{Any, Nothing}, key::CuArray{Float32, 1})
    @ Base .\dict.jl:310
  [8] setindex!(h::Dict{Any, Nothing}, v0::Nothing, key::CuArray{Float32, 1})
    @ Base .\dict.jl:383
  [9] push!(s::Set{Any}, x::CuArray{Float32, 1})
    @ Base .\set.jl:57
 [10] union!(s::Set{Any}, itr::Params)
    @ Base .\abstractset.jl:91
 [11] Set
    @ .\set.jl:10 [inlined]
 [12] _Set
    @ .\set.jl:23 [inlined]
 [13] Set
    @ .\set.jl:21 [inlined]
 [14] issetequal(l::Params, r::Params)
    @ Base .\abstractset.jl:375
 [15] #65
    @ .\none:0 [inlined]
 [16] iterate
    @ .\generator.jl:47 [inlined]
 [17] _all
    @ .\reduce.jl:922 [inlined]
 [18] all
    @ .\reduce.jl:918 [inlined]
 [19] all
    @ .\reduce.jl:836 [inlined]
 [20] map!(::Function, ::Zygote.Grads, ::Zygote.Grads, ::Zygote.Grads)
    @ Zygote C:\Users\tonys\.julia\dev\Zygote\src\compiler\interface.jl:232
 [21] map
    @ C:\Users\tonys\.julia\dev\Zygote\src\compiler\interface.jl:228 [inlined]
 [22] broadcasted(f::Function, gs::Zygote.Grads, gss::Zygote.Grads)
    @ Zygote C:\Users\tonys\.julia\dev\Zygote\src\compiler\interface.jl:211
 [23] top-level scope
    @ REPL[10]:1
 [24] top-level scope
    @ C:\Users\tonys\.julia\packages\CUDA\k52QH\src\initialization.jl:81

[DhairyaLGandhi]

Generally, it is better to use grads such that you extract the gradients explicitly. That is to say, doing g1[p1] .+ g2[p1] would be better.

[darsnack]

Looks like the problem is the check to make sure the Grads being operated on have the same parameter set. We'll have to replace issetequal with a more GPU friendly function.

[DhairyaLGandhi]

We can't guarantee the order in Grads anyway, so I'm not sure if this is the right idiom to support. It's less error prone to explicitly extract the gradient, operate on it, and use that downstream. Probably better yet use a Zygote.hook.

[darsnack]

We can't guarantee the order in Grads anyway, so I'm not sure if this is the right idiom to support.

That was discussed in the original PR. The implementation does not rely on ordering. This is bug that can be fixed, and gradient algebra with many parameters gets verbose.

[DhairyaLGandhi]

We can always extend it to multiple parameters. I'd rather stick with a cleaner implementation.