regression in RNN with OneHotMatrix and CUDA

[CarloLucibello]

The following code is from a test in test/cuda/curnn.jl that used to pass, but I had to comment it out in #1258 since no longer working when CUDA.allowscalar(false)

julia> using Flux

julia> rnn = RNN(10, 5)
Recur(RNNCell(10, 5, tanh))

julia> curnn = rnn |> gpu
Recur(RNNCell(10, 5, tanh))

julia> batch_size = 5  # BATCH_SIZE=1 WORKS FINE!
5

julia> ohx = batch_size == 1 ?
             Flux.onehot(rand(1:10), 1:10) :
             Flux.onehotbatch(rand(1:10, batch_size), 1:10)
10×5 Flux.OneHotMatrix{Array{Flux.OneHotVector,1}}:
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  1  0
 0  0  0  0  0
 1  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  1
 0  1  1  0  0

julia> cuohx = gpu(ohx)
10×5 Flux.OneHotMatrix{CuArray{Flux.OneHotVector,1,Nothing}}:
┌ Warning: Performing scalar operations on GPU arrays: This is very slow, consider disallowing these operations with `allowscalar(false)`
└ @ GPUArrays ~/.julia/packages/GPUArrays/X4SqE/src/host/indexing.jl:43
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  1  0
 0  0  0  0  0
 1  0  0  0  0
 0  0  0  0  0
 0  0  0  0  0
 0  0  0  0  1
 0  1  1  0  0

julia> y = (rnn(ohx); rnn(ohx))
5×5 Array{Float32,2}:
 -0.971044  -0.253226   -0.253226   -0.956147    -0.821044
 -0.973713  -0.528066   -0.528066   -0.968011    -0.71692
  0.646869  -0.263039   -0.263039   -0.00965881   0.349917
  0.307331   0.0366511   0.0366511   0.69206      0.184957
  0.161357  -0.415907   -0.415907    0.0734643    0.279407

julia> cuy = (curnn(cuohx); curnn(cuohx))
5×5 CuArray{Float32,2,Nothing}:
 -0.971044  -0.253226   -0.253226   -0.956147    -0.821045
 -0.973713  -0.528066   -0.528066   -0.968011    -0.71692
  0.646869  -0.263039   -0.263039   -0.00965875   0.349917
  0.307331   0.0366511   0.0366511   0.69206      0.184957
  0.161357  -0.415907   -0.415907    0.0734643    0.279407

julia> CUDA.allowscalar(false)

julia> cuy = (curnn(cuohx); curnn(cuohx))
ERROR: scalar getindex is disallowed
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] assertscalar(::String) at /home/carlo/.julia/packages/GPUArrays/X4SqE/src/host/indexing.jl:41
 [3] getindex(::CuArray{Flux.OneHotVector,1,Nothing}, ::Int64) at /home/carlo/.julia/packages/GPUArrays/X4SqE/src/host/indexing.jl:96
 [4] getindex at /home/carlo/.julia/dev/Flux/src/onehot.jl:23 [inlined]
 [5] _getindex at ./abstractarray.jl:1020 [inlined]
 [6] getindex at ./abstractarray.jl:980 [inlined]
 [7] copyto!(::Array{Float32,2}, ::Flux.OneHotMatrix{CuArray{Flux.OneHotVector,1,Nothing}}) at ./multidimensional.jl:962
 [8] Array at ./array.jl:541 [inlined]
 [9] Array at ./boot.jl:430 [inlined]
 [10] convert at ./array.jl:533 [inlined]
 [11] CuArray at /home/carlo/.julia/packages/CUDA/42B9G/src/array.jl:209 [inlined]
 [12] CuArray at /home/carlo/.julia/packages/CUDA/42B9G/src/array.jl:214 [inlined]
 [13] (::Flux.RNNCell{typeof(tanh),CuArray{Float32,2,Nothing},CuArray{Float32,1,Nothing}})(::CuArray{Float32,2,Nothing}, ::Flux.OneHotMatrix{CuArray{Flux.OneHotVector,1,Nothing}}) at /home/carlo/.julia/dev/Flux/src/cuda/curnn.jl:43
 [14] (::Flux.Recur{Flux.RNNCell{typeof(tanh),CuArray{Float32,2,Nothing},CuArray{Float32,1,Nothing}}})(::Flux.OneHotMatrix{CuArray{Flux.OneHotVector,1,Nothing}}) at /home/carlo/.julia/dev/Flux/src/layers/recurrent.jl:36
 [15] top-level scope at REPL[17]:1
 [16] eval(::Module, ::Any) at ./boot.jl:331
 [17] eval_user_input(::Any, ::REPL.REPLBackend) at /home/carlo/julia/julia-1.4.1/share/julia/stdlib/v1.4/REPL/src/REPL.jl:86
 [18] run_backend(::REPL.REPLBackend) at /home/carlo/.julia/packages/Revise/BqeJF/src/Revise.jl:1184
 [19] top-level scope at none:0

[ToucheSir]

This seems to work as of 0.12.8.