fix gpu tests

Project.toml

@@ -1,11 +1,12 @@
 name = "Flux"
 uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-version = "0.15-DEV"
+version = "0.15.0-DEV"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
+ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
@@ -46,6 +47,7 @@ Adapt = "4"
 CUDA = "5"
 ChainRulesCore = "1.12"
 Compat = "4.10.0"
+ConstructionBase = "1.5.8"
 Enzyme = "0.13"
 Functors = "0.5"
 MLDataDevices = "1.4.2"

src/functor.jl

@@ -108,6 +108,10 @@ julia> m.bias
 """
 cpu(x) = cpu_device()(x)
 
+# TODO remove after https://github.com/LuxDL/Lux.jl/pull/1089
+ChainRulesCore.@non_differentiable cpu_device()
+
+
 # Remove when 
 # https://github.com/JuliaPackaging/Preferences.jl/issues/39
 # is resolved
@@ -149,6 +153,10 @@ CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
 """
 gpu(x) = gpu_device()(x)
 
+# TODO remove after https://github.com/LuxDL/Lux.jl/pull/1089
+ChainRulesCore.@non_differentiable gpu_device()
+ChainRulesCore.@non_differentiable gpu_device(::Any)
+
 # Precision
 
 struct FluxEltypeAdaptor{T} end

src/layers/upsample.jl

@@ -35,6 +35,9 @@ struct Upsample{mode, S, T}
   size::T
 end
 
+Functors.@leaf Upsample # mark leaf since the constructor is not compatible with Functors
+                        # by default but we don't need to recurse into it   
+
 function Upsample(mode::Symbol = :nearest; scale = nothing, size = nothing)
   mode in [:nearest, :bilinear, :trilinear] || 
     throw(ArgumentError("mode=:$mode is not supported."))

test/ext_cuda/layers.jl

@@ -6,11 +6,10 @@
 
 # generic movement tests
 @testset "Basic GPU Movement" begin
-  @test gradient(x -> sum(gpu(x)), rand(3,3)) isa Tuple
-  @test gradient(x -> sum(cpu(x)), gpu(rand(3,3))) isa Tuple
+  @test gradient(x -> sum(gpu(x)), rand(Float32, 3, 3))[1] isa Matrix{Float32}
+  @test gradient(x -> sum(cpu(x)), gpu(rand(Float32, 3, 3)))[1] isa CuMatrix{Float32}
 end
 
-
 const ACTIVATIONS = [identity, tanh]
 
 function gpu_gradtest(name::String, layers::Vector, x_cpu, args...;