copyto method using Adapt move to CPU

Author: mcabbott

src/array.jl

@@ -79,14 +79,14 @@ Base.getindex(x::OneHotArray{<:Any, N}, ::Colon, ::Vararg{Colon, N}) where N = x
 Base.similar(x::OneHotArray{<:Any,<:Any,<:Any,<:AbstractArray}, ::Type{T}, size::Base.Dims) where T =
   similar(x.indices, T, size)
 
-function Base.copyto!(dst::AbstractArray{T,N}, src::OneHotArray{<:Any,Nm1,N,<:AbstractArray}) where {T,N,Nm1}
+function Base.copyto!(dst::AbstractArray{T,N}, src::OneHotArray{<:Any,<:Any,N,<:AbstractArray}) where {T,N}
   size(dst) == size(src) || return invoke(copyto!, Tuple{typeof(dst), AbstractArray{Bool,N}})
-  # fill!(dst, false)
-  # setindex!.(eachslice(dst; dims=ntuple(d->d+1, Nm1)), true, src.indices)
-  # setindex!.(Ref(dst), true, src.indices, axes(src.indices)...)
-  dst .= reshape(src.indices, 1, size(src.indices)...) .== (1:src.nlabels)  # this works at REPL!
+  dst .= reshape(src.indices, 1, size(src.indices)...) .== (1:src.nlabels)
   return dst
 end
+function Base.copyto!(dst::Array{T,N}, src::OneHotArray{<:Any,<:Any,N,<:AnyGPUArray}) where {T,N}
+  copyto!(dst, adapt(Array, src))
+end
 
 function Base.showarg(io::IO, x::OneHotArray, toplevel)
   print(io, ndims(x) == 1 ? "OneHotVector(" : ndims(x) == 2 ? "OneHotMatrix(" : "OneHotArray(")

test/gpu.jl

@@ -37,10 +37,11 @@ end
 
   # These were broken on OneHotArrays v0.2.7
   @test @allowscalar cx[2,2] == x[2,2]
-  @test_broken collect(cx) == collect(x)
-  @test_broken Matrix(cx) == Matrix(x) == collect(x)
-  @test_broken Array{Float32}(cx) == Array{Float32}(x) == collect(x)
+  @test collect(cx) == collect(x)
+  @test Matrix(cx) == Matrix(x) == collect(x)
+  @test Array{Float32}(cx) == Array{Float32}(x) == collect(x)
   @test convert(AbstractArray{Float32}, cx) isa CuArray{Float32}
+  @test collect(convert(AbstractArray{Float32}, cx)) == collect(x)
 end
 
 @testset "onehot gpu" begin