print/convert batchedadjtrans over cuarray (#54)

* print/convert batchedadjtrans over cuarray

* Update test/batchedadjtrans.jl

Co-authored-by: Brian Chen <[email protected]>

Co-authored-by: Brian Chen <[email protected]>

Author: chengchingwen

ext/NNlibCUDA/Project.toml

@@ -3,13 +3,15 @@ uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 version = "0.2.3"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+Adapt = "3.3"
 CUDA = "3.11"
 NNlib = "0.8.7"
 julia = "1.6"

ext/NNlibCUDA/src/NNlibCUDA.jl

@@ -9,6 +9,7 @@ const IntOrIntTuple = Union{Integer, NTuple{N,<:Integer} where N}
 include("upsample.jl")
 include("sampling.jl")
 include("activations.jl")
+include("batchedadjtrans.jl")
 include("batchedmul.jl")
 include("scatter.jl")
 include("gather.jl")

ext/NNlibCUDA/src/batchedadjtrans.jl

@@ -0,0 +1,17 @@
+using NNlib: BatchedAdjoint, BatchedTranspose, BatchedAdjOrTrans
+using Adapt
+using Adapt: WrappedArray
+
+const CuBatchedAdjoint{T} = BatchedAdjoint{T, <: CuArray{T}}
+const CuBatchedTranspose{T} = BatchedTranspose{T, <: CuArray{T}}
+const CuBatchedAdjOrTrans{T} = Union{CuBatchedAdjoint{T}, CuBatchedTranspose{T}}
+const WrappedCuBatchedAdjOrTrans{T, N} = WrappedArray{T, N, CuBatchedAdjOrTrans{T}, CuBatchedAdjOrTrans{T}}
+
+
+Base.print_array(io::IO, b::Union{CuBatchedAdjOrTrans, WrappedCuBatchedAdjOrTrans}) = Base.print_array(io, adapt(Array, b))
+Base._show_nonempty(io::IO, b::Union{CuBatchedAdjOrTrans, WrappedCuBatchedAdjOrTrans}, prefix::String) = Base._show_nonempty(io, adapt(Array, b), prefix)
+Base.show_vector(io::IO, b::Union{CuBatchedAdjOrTrans, WrappedCuBatchedAdjOrTrans}, opn, cls) = Base.show_vector(io, adapt(Array, b), opn, cls)
+
+Base.convert(::Type{T}, b::Union{CuBatchedAdjOrTrans, WrappedCuBatchedAdjOrTrans}) where {T<:Array} = Base.convert(T, adapt(Array, b))
+Base.Array{T, N}(b::Union{CuBatchedAdjOrTrans, WrappedCuBatchedAdjOrTrans}) where {T, N} = Array{T, N}(adapt(Array, b))
+Base.collect(b::Union{CuBatchedAdjOrTrans, WrappedCuBatchedAdjOrTrans}) = collect(adapt(Array, b))

ext/NNlibCUDA/test/batchedadjtrans.jl

@@ -0,0 +1,44 @@
+function print_array_strs(x)
+    str = sprint((io, x)->show(io, MIME"text/plain"(), x), x)
+    return @view split(str, '\n')[2:end]
+end
+
+@testset "BatchedAdjOrTrans" begin
+    x = randn(Float32, 3,4,2)
+    y = cu(x)
+
+    bax = batched_adjoint(x)
+    btx = batched_transpose(x)
+    bay = batched_adjoint(y)
+    bty = batched_transpose(y)
+
+    @test sprint(show, bax) == sprint(show, bay)
+    @test sprint(show, btx) == sprint(show, bty)
+
+    @test print_array_strs(bax) == print_array_strs(bay)
+    @test print_array_strs(btx) == print_array_strs(bty)
+
+    @test Array(bax) == Array(bay)
+    @test collect(bax) == collect(bay)
+    @test Array(btx) == Array(bty)
+    @test collect(btx) == collect(bty)
+    
+    for shape in (:, (12, 2))
+        rbax = reshape(bax, shape)
+        rbtx = reshape(btx, shape)
+        rbay = reshape(bay, shape)
+        rbty = reshape(bty, shape)
+
+        @test sprint(show, rbax) == sprint(show, rbay)
+        @test sprint(show, rbtx) == sprint(show, rbty)
+    
+        @test print_array_strs(rbax) == print_array_strs(rbay)
+        @test print_array_strs(rbtx) == print_array_strs(rbty)
+    
+        @test Array(rbax) == Array(rbay)
+        @test collect(rbax) == collect(rbay)
+        @test Array(rbtx) == Array(rbty)
+        @test collect(rbtx) == collect(rbty)
+    end
+
+end

ext/NNlibCUDA/test/runtests.jl

@@ -10,6 +10,7 @@ CUDA.allowscalar(false)
 @testset "NNlibCUDA" begin
 include("test_utils.jl")
 include("activations.jl")
+include("batchedadjtrans.jl")
 include("batchedmul.jl")
 include("upsample.jl")
 include("conv.jl")