diff --git a/Project.toml b/Project.toml index 70673996d7..e36fe05370 100644 --- a/Project.toml +++ b/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" diff --git a/src/functor.jl b/src/functor.jl index 5a97836cfa..e8c02b919f 100644 --- a/src/functor.jl +++ b/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 diff --git a/src/layers/upsample.jl b/src/layers/upsample.jl index c71a9acc8d..5f79b98fde 100644 --- a/src/layers/upsample.jl +++ b/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.")) diff --git a/test/ext_cuda/layers.jl b/test/ext_cuda/layers.jl index cdb8f003e9..b7b456bcf1 100644 --- a/test/ext_cuda/layers.jl +++ b/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...;