Inline printing for OneHotArray is not GPU-friendly #1905
Description
Found in https://discourse.julialang.org/t/scalar-indexing-error-when-loading-mnist/77496:
julia> Flux.onehotbatch(1:2, 1:2) |> cpu, Flux.onehotbatch(1:2, 1:2) |> cpu
(Bool[1 0; 0 1], Bool[1 0; 0 1])
julia> Flux.onehotbatch(1:2, 1:2) |> gpu, Flux.onehotbatch(1:2, 1:2) |> gpu
(Bool[Error showing value of type Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}}:
ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore are only permitted from the REPL for prototyping purposes.
If you did intend to index this array, annotate the caller with @allowscalar.
Stacktrace:
[1] error(s::String)
@ Base ./error.jl:33
[2] assertscalar(op::String)
@ GPUArrays ~/.julia/packages/GPUArrays/gkF6S/src/host/indexing.jl:53
[3] getindex
@ ~/.julia/packages/GPUArrays/gkF6S/src/host/indexing.jl:86 [inlined]
[4] getindex
@ ~/.julia/packages/Flux/BPPNj/src/onehot.jl:46 [inlined]
[5] isassigned(::Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, ::Int64, ::Int64)
@ Base ./abstractarray.jl:553
[6] _show_nonempty(io::IOContext{Base.TTY}, X::AbstractMatrix, prefix::String, drop_brackets::Bool, axs::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}})
@ Base ./arrayshow.jl:438
[7] _show_nonempty(io::IOContext{Base.TTY}, X::Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, prefix::String)
@ Base ./arrayshow.jl:410
[8] show
@ ./arrayshow.jl:486 [inlined]
[9] show_delim_array(io::IOContext{Base.TTY}, itr::Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}}, op::Char, delim::Char, cl::Char, delim_one::Bool, i1::Int64, n::Int64)
@ Base ./show.jl:1244
[10] show_delim_array
@ ./show.jl:1229 [inlined]
[11] show
@ ./show.jl:1262 [inlined]
[12] show(io::IOContext{Base.TTY}, #unused#::MIME{Symbol("text/plain")}, x::Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}})
@ Base.Multimedia ./multimedia.jl:47
[13] (::REPL.var"#43#44"{REPL.REPLDisplay{REPL.LineEditREPL}, MIME{Symbol("text/plain")}, Base.RefValue{Any}})(io::Any)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:264
[14] with_repl_linfo(f::Any, repl::REPL.LineEditREPL)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:508
[15] display(d::REPL.REPLDisplay, mime::MIME{Symbol("text/plain")}, x::Any)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:257
[16] display(d::REPL.REPLDisplay, x::Any)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:269
[17] display(x::Any)
@ Base.Multimedia ./multimedia.jl:328
[18] (::Media.var"#15#16"{Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}}})()
@ Media ~/.julia/packages/Media/ItEPc/src/compat.jl:28
[19] hookless(f::Media.var"#15#16"{Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}}})
@ Media ~/.julia/packages/Media/ItEPc/src/compat.jl:14
[20] render(#unused#::Media.NoDisplay, x::Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}})
@ Media ~/.julia/packages/Media/ItEPc/src/compat.jl:27
[21] render(x::Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}})
@ Media ~/.julia/packages/Media/ItEPc/src/system.jl:160
[22] display(#unused#::Media.DisplayHook, x::Tuple{Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}, Flux.OneHotArray{UInt32, 2, 1, 2, CuArray{UInt32, 1, CUDA.Mem.DeviceBuffer}}})
@ Media ~/.julia/packages/Media/ItEPc/src/compat.jl:9
[23] display(x::Any)
@ Base.Multimedia ./multimedia.jl:328
[24] #invokelatest#2
@ ./essentials.jl:716 [inlined]
[25] invokelatest
@ ./essentials.jl:714 [inlined]
[26] print_response(errio::IO, response::Any, show_value::Bool, have_color::Bool, specialdisplay::Union{Nothing, AbstractDisplay})
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:291
[27] (::REPL.var"#45#46"{REPL.LineEditREPL, Pair{Any, Bool}, Bool, Bool})(io::Any)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:275
[28] with_repl_linfo(f::Any, repl::REPL.LineEditREPL)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:508
[29] print_response(repl::REPL.AbstractREPL, response::Any, show_value::Bool, have_color::Bool)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:273
[30] (::REPL.var"#do_respond#66"{Bool, Bool, REPL.var"#77#87"{REPL.LineEditREPL, REPL.REPLHistoryProvider}, REPL.LineEditREPL, REPL.LineEdit.Prompt})(s::REPL.LineEdit.MIState, buf::Any, ok::Bool)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:844
[31] #invokelatest#2
@ ./essentials.jl:716 [inlined]
[32] invokelatest
@ ./essentials.jl:714 [inlined]
[33] run_interface(terminal::REPL.Terminals.TextTerminal, m::REPL.LineEdit.ModalInterface, s::REPL.LineEdit.MIState)
@ REPL.LineEdit ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/LineEdit.jl:2493
[34] run_frontend(repl::REPL.LineEditREPL, backend::REPL.REPLBackendRef)
@ REPL ~/.julia/juliaup/julia-1.7.1+0~x64/share/julia/stdlib/v1.7/REPL/src/REPL.jl:1230
[35] (::REPL.var"#49#54"{REPL.LineEditREPL, REPL.REPLBackendRef})()
@ REPL ./task.jl:423
Looking at the stacktrace, it appears none of https://github.com/FluxML/Flux.jl/blob/master/src/onehot.jl#L51-L68 is being hit.