v0.3.7

Project.toml

@@ -1,7 +1,7 @@
 name = "FourierFilterFlux"
 uuid = "3d7dfd45-6c90-4c9b-b697-194a05757159"
 authors = ["dsweber2"]
-version = "0.3.5"
+version = "0.3.7"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/FourierFilterFlux.jl

@@ -1,13 +1,15 @@
 module FourierFilterFlux
 using Reexport
-# @reexport using CUDA
-using CUDA
 using Zygote, Flux, Adapt, LinearAlgebra
 using AbstractFFTs, FFTW # TODO: check the license on FFTW and such
 using ContinuousWavelets
 using RecipesBase
+using CUDA
 
 const use_cuda = Ref(false)
+if CUDA.functional()
+    use_cuda[] = true
+end
 
 import Adapt: adapt
 export pad, originalDomain, formatJLD, getBatchSize

src/Utils.jl

@@ -1,7 +1,7 @@
 import NNlib.relu
 # just a little bit of type piracy used internally TODO maybe don't...
 relu(x::C) where {C<:Complex} = real(x) > 0 ? x : C(0)
-
+import CUDA: CuArray
 # ways to convert between gpu and cpu
 import Adapt.adapt
 function adapt(to, cft::ConvFFT{D,OT,F,A,V,PD,P,T,An}) where {D,OT,F,A,V,PD,P,T,An}
@@ -38,16 +38,23 @@ function cu(cft::ConvFFT{D,OT,F,A,V,PD,P,T,An}) where {D,OT,F,A,V,PD,P,T,An}
         cft.analytic)
 end
 
-# TODO this is somewhat kludgy, not sure why cu was converting these back
-#function CUDA.cu(P::FFTW.rFFTWPlan)
-#    return plan_rfft(cu(zeros(real(eltype(P)), P.sz)), P.region)
-#end
-#CUDA.cu(P::CUFFT.rCuFFTPlan) = P
+# Only convert FFTW plans to CUFFT plans if CUDA is actually functional
+function CUDA.cu(P::FFTW.rFFTWPlan)
+    if CUDA.functional()
+        return plan_rfft(CUDA.cu(zeros(real(eltype(P)), P.sz)), P.region)
+    else
+        return P   # fallback to CPU FFTW plan
+    end
+end
 
-#function CUDA.cu(P::FFTW.cFFTWPlan)
-#    return plan_fft(cu(zeros(eltype(P), P.sz)), P.region)
-#end
-#CUDA.cu(P::CUFFT.cCuFFTPlan) = P
+function CUDA.cu(P::FFTW.cFFTWPlan)
+    if CUDA.functional()
+        return plan_fft(CUDA.cu(zeros(eltype(P), P.sz)), P.region)
+    else
+        return P   # fallback to CPU FFTW plan
+    end
+end
+CUDA.cu(P::CUDA.CUFFT.Plan) = P
 
 Adapt.adapt(::Type{Array{T}}, P::FFTW.FFTWPlan{T}) where {T} = P
 function Adapt.adapt(::Type{Array{T}}, P::FFTW.rFFTWPlan) where {T}
@@ -62,8 +69,8 @@ adapt(::Type{<:CuArray}, x::T) where {T<:CUDA.CUFFT.CuFFTPlan} = x
 # is actually converting
 function adapt(::Union{Type{<:Array},Flux.FluxCPUAdaptor},
         x::T) where {T<:CUDA.CUFFT.CuFFTPlan}
-    transformSize = x.osz
-    dataSize = x.sz
+    transformSize = x.output_size
+    dataSize = x.input_size
     if dataSize != transformSize
         # this is an rfft, since the dimension isn't preserved
         return plan_rfft(zeros(dataSize), x.region)

src/boundaries.jl

@@ -11,7 +11,7 @@ end
 
 N gives the number of dimensions of convolution, while `x` gives the specific amount to pad in each dimension (done on both sides). If the values in `x` are negative, then the support of the filters will be determined automataically
 """
-Pad(x::Vararg{<:Integer,N}) where {N} = Pad{N}(x)
+Pad(x::Vararg{Integer,N}) where {N} = Pad{N}(x)
 import Base.ndims
 ndims(p::Pad{N}) where {N} = N
 

src/convFFTConstructors.jl

@@ -78,7 +78,7 @@ function waveletLayer(inputSize::Union{T,NTuple{N,T}};
             averagingStyle = RealWaveletComplexSignal
         end
         An = map(ii -> ((ii in An) ? averagingStyle() :
-                        AnalyticWavelet()), (1:size(wavelets, 2)[end]...,))
+                    AnalyticWavelet()), 1:size(wavelets, 2))
     end
     if bias
         bias = dType.(init(inputSize[2:end-1]..., size(wavelets, 2)))

src/paramCollection.jl

@@ -1,4 +1,4 @@
-Flux.@functor ConvFFT
+Flux.@layer ConvFFT
 
 function Flux.trainable(CFT::ConvFFT{A,B,C,D,E,F,G,true}) where {A,B,C,D,E,F,G}
     (CFT.weight, CFT.bias)

src/transforms.jl

@@ -1,5 +1,5 @@
 # TODO: version that doesn't have an fft built in
-
+import CUDA: CuArray
 function (shears::ConvFFT)(x)
     if typeof(shears.weight) <: CuArray && !(typeof(x) <: CuArray)
         error("don't try to apply a gpu transform to a non-CuArray")

test/CUDATests.jl

@@ -1,13 +1,13 @@
 if CUDA.functional()
     @testset "CUDA methods" begin
         w = ConvFFT((100,), nConvDims = 1)
-        @test cu(w.fftPlan) isa CUFFT.rCuFFTPlan # does cu work on the fft plans when applied directly?
+        @test cu(w.fftPlan) isa CUDA.CUFFT.CuFFTPlan # does cu work on the fft plans when applied directly?
         cw = cu(w)
         @test cw.weight isa NTuple{N,CuArray} where {N} # does cu work on the weights?
-        @test cw.fftPlan isa CUFFT.rCuFFTPlan # does cu work on the fftPlan?
+        @test cw.fftPlan isa CUDA.CUFFT.CuFFTPlan # does cu work on the fftPlan?
         cw1 = gpu(w)
         @test cw1.weight isa NTuple{N,CuArray} where {N} # does gpu work on the weights?
-        @test cw1.fftPlan isa CUFFT.rCuFFTPlan # does gpu work on the fftPlan?
+        @test cw1.fftPlan isa CUDA.CUFFT.CuFFTPlan # does gpu work on the fftPlan?
         w1 = cpu(cw)
         @test w1.weight isa NTuple{N,Array} where {N} # does cpu work on the weights?
         @test w1.fftPlan isa FFTW.rFFTWPlan # does cpu work on the fftPlan?
@@ -16,15 +16,15 @@ if CUDA.functional()
         @test cw(cx) isa CuArray
         @test cw(cx) ≈ cu(w(x)) # CUDA and cpu version get the same result approximately
         cw(cx)
-        ∇cu = gradient(t -> sum(cw(t)), cx)[1]
-        ∇ = gradient(t -> sum(w(t)), x)[1]
+        ∇cu = Flux.gradient(t -> sum(cw(t)), cx)[1]
+        ∇ = Flux.gradient(t -> sum(w(t)), x)[1]
         @test ∇ ≈ cpu(∇cu)
         w1 = waveletLayer((100, 1, 1))
         cw1 = cu(w1)
         @test cw1(cx) ≈ cu(w1(x))
 
-        CUDA.@allowscalar ∇cu = gradient(t -> abs(cw1(t)[1]), cx)[1]
-        CUDA.@allowscalar ∇ = gradient(t -> abs(w1(t)[1]), x)[1]
+        CUDA.@allowscalar ∇cu = Flux.gradient(t -> abs(cw1(t)[1]), cx)[1]
+        CUDA.@allowscalar ∇ = Flux.gradient(t -> abs(w1(t)[1]), x)[1]
         @test ∇ ≈ cpu(∇cu)
     end
 end

test/ConvFFTConstructors.jl

@@ -1,6 +1,7 @@
 # TODO: add some checks for different boundary conditions
 # TODO: add checks for analytic wavelets
 # ConvFFT constructor tests
+using FourierFilterFlux: applyWeight, applyBC, internalConvFFT
 @testset "ConvFFT constructors" begin
     @testset "Utils" begin
         explicit = [1 0 0; 0 1 0; 0 0 1; zeros(2, 3)]
@@ -22,16 +23,17 @@
 
         shears = ConvFFT(weightMatrix, nothing, originalSize, abs,
             plan = true, boundary = Pad(padding), trainable = true)
-        @test Flux.params(shears).order[1] == shears.weight[1]
-        @test length(Flux.params(shears).order) == 1
+         trn = Flux.trainable(shears)
+        @test trn[1][1] == shears.weight[1]
+        @test length(trn) >= 1
 
         shears = ConvFFT(weightMatrix, nothing, originalSize, abs,
             boundary = Pad(padding), trainable = false)
-        @test isempty(Flux.params(shears))
+        @test isempty(Flux.trainable(shears))
 
         x = randn(21, 11, 1, 10)
-        ∇ = gradient((x) -> shears(x)[1, 1, 1, 1, 3], x)
-        @test minimum(∇[1][:, :, :, [1:2..., 4:10...]] .≈ 0)
+        ∇ = Flux.gradient((x) -> shears(x)[1, 1, 1, 1, 3], x)
+        @test all(∇[1][:, :, :, [1:2..., 4:10...]] .≈ 0)
 
         # check that the identity ConvFFT is, in fact, an identity
         weightMatrix = ones(Float32, (21 + 10) >> 1 + 1, 11 + 10, 1)
@@ -59,7 +61,7 @@
         nextLayer = FourierFilterFlux.internalConvFFT(x̂, shears.weight, usedInds,
             shears.fftPlan, shears.bias,
             shears.analytic)
-        ∇ = gradient((x̂) -> FourierFilterFlux.internalConvFFT(x̂,
+        ∇ = Flux.gradient((x̂) -> FourierFilterFlux.internalConvFFT(x̂,
                 shears.weight,
                 usedInds,
                 shears.fftPlan,
@@ -69,20 +71,20 @@
                 1,
                 1],
             x̂)
-        @test minimum(abs.(diag(∇[1][:, :, 1, 1])) .≈ 2.0f0 / 31 / 21)
+        @test all(abs.(diag(∇[1][:, :, 1, 1])) .≈ 2.0f0 / 31 / 21)
 
         ax = axes(x̂)[3:end-1]
-        ∇ = gradient((x̂) -> FourierFilterFlux.applyWeight(x̂, shears.weight[1], usedInds,
+        ∇ = Flux.gradient((x̂) -> FourierFilterFlux.applyWeight(x̂, shears.weight[1], usedInds,
                 shears.fftPlan,
                 shears.bias, FourierFilterFlux.NonAnalyticMatching())[1,
                 1,
                 1,
                 1,
                 1], x̂)
-        @test minimum(abs.(diag(∇[1][:, :, 1, 1])) .≈ 2.0f0 / 31 / 21)
+        @test all(abs.(diag(∇[1][:, :, 1, 1])) .≈ 2.0f0 / 31 / 21)
 
-        ∇ = gradient((x̂) -> (shears.fftPlan\(x̂.*shears.weight[1]))[1, 1, 1, 1], x̂)
-        @test minimum(abs.(diag(∇[1][:, :, 1, 1])) .≈ 1.0f0 / 31 * 2 / 21)
+        ∇ = Flux.gradient((x̂) -> (shears.fftPlan\(x̂.*shears.weight[1]))[1, 1, 1, 1], x̂)
+        @test all(abs.(diag(∇[1][:, :, 1, 1])) .≈ 1.0f0 / 31 * 2 / 21)
         sheared = shears(x)
         @test size(sheared) == (21, 11, 1, 1, 10)
 
@@ -97,11 +99,11 @@
         if CUDA.functional()
             gpuVer = shears |> gpu
             @test gpuVer.weight[1] isa CuArray
-            @test gpuVer.fftPlan isa CUFFT.rCuFFTPlan
+            @test gpuVer.fftPlan isa CUDA.CUFFT.CuFFTPlan
             if !(gpuVer.weight[1] isa CuArray)
                 println("gpuVer.weight is of type $(typeof(gpuVer.weight))")
             end
-            if !(gpuVer.fftPlan isa CUFFT.rCuFFTPlan)
+            if !(gpuVer.fftPlan isa CUDA.CUFFT.CuFFTPlan)
                 println("gpuVer.fftPlan is of type $(typeof(gpuVer.fftPlan))")
             end
         end
@@ -147,15 +149,17 @@
             @test shears.σ == abs
             @test shears.bias == nothing
             @test shears.bc.padBy == (5,)
-            @test Flux.params(shears).order[1] == shears.weight[1]
+            trn = Flux.trainable(shears)
+            @test trn[1][1] == shears.weight[1] 
+            @test length(trn) >= 1
 
             shears = ConvFFT(weightMatrix, nothing, originalSize, abs,
                 plan = true, boundary = Pad(padding), trainable = false)
-            @test isempty(Flux.params(shears))
+            @test isempty(Flux.trainable(shears))
 
             x = randn(21, 1, 10)
-            ∇ = gradient((x) -> shears(x)[1, 1, 1, 3], x)
-            @test minimum(∇[1][:, :, [1:2..., 4:10...]] .≈ 0)
+            ∇ = Flux.gradient((x) -> shears(x)[1, 1, 1, 3], x)
+            @test all(∇[1][:, :, [1:2..., 4:10...]] .≈ 0)
 
             # Sym test
             weightMatrix = randn(Float32, (21 + 1), 1)
@@ -165,23 +169,27 @@
             @test shears.σ == abs
             @test shears.bias == nothing
             @test typeof(shears.bc) <: Sym
-            @test Flux.params(shears).order[1] == shears.weight[1]
+            trn = Flux.trainable(shears)
+            @test trn[1][1] == shears.weight[1]
+            @test length(trn) >= 1
             x = randn(21, 1, 10)
-            ∇ = gradient((x) -> shears(x)[1, 1, 1, 3], x)
-            @test minimum(∇[1][:, :, [1:2..., 4:10...]] .≈ 0)
-            @test minimum(abs.(∇[1][:, 1, 3])) > 0
+            ∇ = Flux.gradient((x) -> shears(x)[1, 1, 1, 3], x)
+            @test all(∇[1][:, :, [1:2..., 4:10...]] .≈ 0)
+            @test all(abs.(∇[1][:, 1, 3]) .> 0)
             weightMatrix = randn(Float32, 21 >> 1 + 1, 1)
             shears = ConvFFT(weightMatrix, nothing, originalSize, abs,
                 plan = true, boundary = FourierFilterFlux.Periodic())
             @test size(shears.fftPlan) == originalSize
             @test shears.σ == abs
             @test shears.bias == nothing
             @test typeof(shears.bc) <: FourierFilterFlux.Periodic
-            @test Flux.params(shears).order[1] == shears.weight[1]
+            trn = Flux.trainable(shears)
+            @test trn[1][1] == shears.weight[1]
+            @test length(trn) >= 1
             x = randn(21, 1, 10)
-            ∇ = gradient((x) -> shears(x)[1, 1, 1, 3], x)
-            @test minimum(∇[1][:, :, [1:2..., 4:10...]] .≈ 0)
-            @test minimum(abs.(∇[1][:, 1, 3])) > 0
+            ∇ = Flux.gradient((x) -> shears(x)[1, 1, 1, 3], x)
+            @test all(∇[1][:, :, [1:2..., 4:10...]] .≈ 0)
+            @test all(abs.(∇[1][:, 1, 3]) .> 0)
         end
 
         # check that the identity ConvFFT is, in fact, an identity
@@ -235,7 +243,6 @@
         end
 
 
-        using FourierFilterFlux: applyWeight, applyBC, internalConvFFT
         weight = (2 .* ones(Complex{Float32}, (21 + 10) >> 1 + 1),)
         bc = Pad(5)
         x = randn(Float32, 21, 1, 10)
@@ -244,13 +251,13 @@
         fftPlan = plan_rfft(xbc, (1,))
         An = map(x -> FourierFilterFlux.NonAnalyticMatching(), (1:length(weight)...,))
         nextLayer = internalConvFFT(x̂, weight, usedInds, fftPlan, nothing, An)
-        ∇ = gradient((x̂) -> internalConvFFT(x̂, weight, usedInds, fftPlan, nothing, An)[1,
+        ∇ = Flux.gradient((x̂) -> internalConvFFT(x̂, weight, usedInds, fftPlan, nothing, An)[1,
                 1,
                 1,
                 1,
                 1],
             x̂)
-        y, ∂ = pullback((x̂) -> internalConvFFT(x̂, weight, usedInds, fftPlan, nothing, An)[1,
+        y, ∂ = Zygote.pullback((x̂) -> internalConvFFT(x̂, weight, usedInds, fftPlan, nothing, An)[1,
                 1,
                 1,
                 1,
@@ -259,12 +266,12 @@
         ∂(y)
         ∂(y) # repeated calls to the derivative were causing errors while argWrapper
         # was in use
-        @test minimum(abs.(∇[1][:, 1, 1]) .≈ 2.0f0 / 31)
+        @test all(abs.(∇[1][:, 1, 1]) .≈ 2.0f0 / 31)
         # no bias, not analytic and real valued output
 
         # no bias, analytic (so complex valued)
         fftPlan = plan_fft(xbc, (1,))
-        ∇ = gradient((x̂) -> abs(applyWeight(x̂,
+        ∇ = Flux.gradient((x̂) -> abs(applyWeight(x̂,
                 weight[1],
                 usedInds,
                 fftPlan,
@@ -274,7 +281,7 @@
                 1,
                 1]),
             x̂)
-        @test minimum(abs.(∇[1][:, 1, 1]) .≈ 2.0f0 / 31)
+        @test all(abs.(∇[1][:, 1, 1]) .≈ 2.0f0 / 31)
 
         # no bias, not analytic, complex valued, but still symmetric
         real(applyWeight(x̂,
@@ -284,7 +291,7 @@
             nothing,
             FourierFilterFlux.RealWaveletRealSignal()))
         fftPlan = plan_fft(xbc, (1,))
-        ∇ = gradient((x̂) -> real(applyWeight(x̂,
+        ∇ = Flux.gradient((x̂) -> real(applyWeight(x̂,
                 weight[1],
                 usedInds,
                 fftPlan,
@@ -294,7 +301,7 @@
                 1,
                 1]),
             x̂)
-        @test minimum(abs.(∇[1][2:end, 1, 1]) .≈ 2 * 2.0f0 / 31)
+        @test all(abs.(∇[1][2:end, 1, 1]) .≈ 2 * 2.0f0 / 31)
         @test abs(∇[1][1, 1, 1]) ≈ 2.0f0 / 31
 
         # internal methods tests
@@ -309,7 +316,7 @@
         nextLayer = FourierFilterFlux.internalConvFFT(x̂, shears.weight, usedInds,
             shears.fftPlan,
             shears.bias, shears.analytic)
-        ∇ = gradient((x̂) -> FourierFilterFlux.internalConvFFT(x̂,
+        ∇ = Flux.gradient((x̂) -> FourierFilterFlux.internalConvFFT(x̂,
                 shears.weight,
                 usedInds,
                 shears.fftPlan,
@@ -319,24 +326,24 @@
                 1,
                 1],
             x̂)
-        @test minimum(abs.(∇[1][:, 1, 1]) .≈ 2.0f0 / 31)
+        @test all(abs.(∇[1][:, 1, 1]) .≈ 2.0f0 / 31)
         #
 
         # no bias, not analytic and real valued output
         # no bias, analytic (so complex valued)
         # no bias, not analytic, complex valued, but still symmetric
         # biased (and one of the others, doesn't matter which)
 
-        ∇ = gradient((x̂) -> (shears.fftPlan\(x̂.*shears.weight[1]))[1, 1, 1, 1], x̂)
-        @test minimum(abs.(∇[1][:, :, 1, 1]) .≈ 1.0f0 / 31 * 2)
+        ∇ = Flux.gradient((x̂) -> (shears.fftPlan\(x̂.*shears.weight[1]))[1, 1, 1, 1], x̂)
+        @test all(abs.(∇[1][:, :, 1, 1]) .≈ 1.0f0 / 31 * 2)
         sheared = shears(x)
         @test size(sheared) == (21, 1, 1, 10)
 
         # convert to a gpu version
         if CUDA.functional()
             gpuVer = shears |> gpu
             @test gpuVer.weight[1] isa CuArray
-            @test gpuVer.fftPlan isa CUFFT.rCuFFTPlan
+            @test gpuVer.fftPlan isa CUDA.CUFFT.CuFFTPlan
         end
         # extra channel dimension
         originalSize = (20, 16, 1, 10)