Merge pull request #2034 from chengchingwen/master

fix array constructor rrule

Author: chengchingwen

src/functor.jl

@@ -115,7 +115,7 @@ adapt_storage(to::FluxCPUAdaptor, x::AbstractSparseArray) = x
 adapt_storage(to::FluxCPUAdaptor, x::CUDA.RNG) = Random.default_rng()
 adapt_storage(to::FluxCPUAdaptor, x::AbstractRNG) = x
 
-function ChainRulesCore.rrule(::typeof(Array), x::CUDA.CuArray)
+function ChainRulesCore.rrule(::Type{Array}, x::CUDA.CuArray)
   Array(x), d -> (NoTangent(), CUDA.cu(d),)
 end
 

test/utils.jl

@@ -765,5 +765,16 @@ end
     g = Flux.Zygote.ForwardDiff.gradient(pv -> loss(data, 1, pv), pvec)
     @test g ≈ Flux.Zygote.gradient(pv -> loss(data, 1, pv), pvec)[1]
   end
+end
 
+@testset "Rrule" begin
+  @testset "issue 2033" begin
+    if CUDA.functional()
+      struct Wrapped{T}
+          x::T
+      end
+      y, _ = Flux.pullback(Wrapped, cu(randn(3,3)))
+      @test y isa Wrapped{<:CuArray}
+    end
+  end
 end