From a7d16ebc4f9e4fb737ae8cb5e77ebb2abe13aa5e Mon Sep 17 00:00:00 2001
From: JaredW40 <trumpetpro40@gmail.com>
Date: Tue, 18 Aug 2026 16:22:56 -0700
Subject: [PATCH 1/5] Updated Project.toml and CI.yml to properly load
 necessary packages for tests.

---
 .github/workflows/CI.yml | 53 +++++++++++++++++++++++++++++-----------
 Project.toml             | 12 ++++++---
 2 files changed, 47 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index c296bb0..91c9d35 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -25,29 +25,24 @@ env:
 
 jobs:
   test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.group }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        version:
-          - '1.12'
-        os:
-          - ubuntu-latest
-          - macOS-latest   # Apple Silicon (ARM64)
-        arch:
-          - x64
-        exclude:
-          - os: macOS-latest
-            arch: x64
         include:
+          - os: ubuntu-latest
+            arch: x64
+            version: '1.12'
+            group: CUDA
           - os: macOS-latest
             arch: aarch64
             version: '1.12'
+            group: Metal
 
     steps:
       - uses: actions/checkout@v4
-      
+
       - name: Check out ContinuousWavelets.jl
         uses: actions/checkout@v4
         with:
@@ -96,6 +91,7 @@ jobs:
             ])
             Pkg.add([
               PackageSpec(name = "Test"),
+              PackageSpec(name = "BenchmarkTools"),
               PackageSpec(name = "Wavelets"),
               PackageSpec(name = "Zygote"),
               PackageSpec(name = "Flux"),
@@ -127,14 +123,43 @@ jobs:
       - name: Precompile
         run: julia --project=ci -e 'using Pkg; Pkg.precompile()'
 
+      - name: Environment status
+        run: julia --project=ci -e 'using Pkg; Pkg.status()'
+
+      - name: Report GPU availability
+        continue-on-error: true
+        run: |
+          julia --project=ci -e '
+            if Sys.islinux()
+                using CUDA
+                @info "CUDA backend" functional = CUDA.functional()
+            elseif Sys.isapple()
+                using Metal
+                @info "Metal backend" functional = Metal.functional()
+            end'
+
       - name: Run tests
         env:
-          GROUP: ${{ runner.os == 'Linux' && 'CUDA' || runner.os == 'macOS' && 'Metal' || 'CPU' }}
+          GROUP: ${{ matrix.group }}
         run: julia --project=ci --code-coverage=user -e 'include("test/runtests.jl")'
 
       - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: >-
+            src,ext,
+            deps/ContinuousWavelets.jl/src,deps/ContinuousWavelets.jl/ext,
+            deps/FourierFilterFlux.jl/src,deps/FourierFilterFlux.jl/ext,
+            deps/MonogenicFilterFlux.jl/src,deps/MonogenicFilterFlux.jl/ext
+
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: lcov-${{ matrix.os }}-${{ matrix.group }}
+          path: lcov.info
+          if-no-files-found: warn
+
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
\ No newline at end of file
+          fail_ci_if_error: false
diff --git a/Project.toml b/Project.toml
index eadf0f9..56a2c4b 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,7 +9,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ContinuousWavelets = "96eb917e-2868-4417-9cb6-27e7ff17528f"
-CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
@@ -25,20 +24,21 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-CUDA_Runtime_Discovery = "1af6417a-86b4-443c-805f-a4643ffb695f"
+CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
 cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 cuFFT = "533571aa-0936-420e-b4be-9c66f5f626ca"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
-CUDAExt = ["CUDA", "cuDNN", "cuFFT"]
+CUDAExt = ["CUDA", "cuDNN", "cuFFT", "CUDNN_jll"]
 MetalExt = "Metal"
 ScatteringPlotsExt = "Plots"
 
 [compat]
 AbstractFFTs = "1"
 Adapt = "4"
+BenchmarkTools = "1"
 ContinuousWavelets = "1"
 CUDA = "6.1"
 cuDNN = "6.1"
@@ -48,13 +48,17 @@ FFTW = "1"
 Flux = "0.13, 0.14, 0.15, 0.16"
 Functors = "0.5.2"
 julia = "1.12"
+Metal = "1"
 Plots = "1.41.6"
+Statistics = "1"
 Wavelets = "0.9, 0.10"
 Zygote = "0.6, 0.7"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "BenchmarkTools"]
+test = ["Test", "BenchmarkTools", "Logging", "Random"]

From bc904c48ac53470f05e3f19e9a7a2f08f99cf2c0 Mon Sep 17 00:00:00 2001
From: JaredW40 <trumpetpro40@gmail.com>
Date: Tue, 18 Aug 2026 16:57:59 -0700
Subject: [PATCH 2/5] NNlib contains no Metal pooling Kernel, so the mean
 pooling functions fail. This is only an issue for Metal. To remedy, we create
 functions to accomplish these goals.

---
 .github/workflows/CI.yml |  8 +---
 ext/MetalExt.jl          | 91 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 91c9d35..577ebca 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -144,12 +144,6 @@ jobs:
         run: julia --project=ci --code-coverage=user -e 'include("test/runtests.jl")'
 
       - uses: julia-actions/julia-processcoverage@v1
-        with:
-          directories: >-
-            src,ext,
-            deps/ContinuousWavelets.jl/src,deps/ContinuousWavelets.jl/ext,
-            deps/FourierFilterFlux.jl/src,deps/FourierFilterFlux.jl/ext,
-            deps/MonogenicFilterFlux.jl/src,deps/MonogenicFilterFlux.jl/ext
 
       - uses: actions/upload-artifact@v4
         if: always()
@@ -162,4 +156,4 @@ jobs:
         with:
           files: lcov.info
           token: ${{ secrets.CODECOV_TOKEN }}
-          fail_ci_if_error: false
+          fail_ci_if_error: false
\ No newline at end of file
diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl
index 42a1bb6..26efede 100644
--- a/ext/MetalExt.jl
+++ b/ext/MetalExt.jl
@@ -2,4 +2,95 @@ module MetalExt
 
 using ScatteringTransform, Metal
 
+#= NNlib ships no Metal pooling kernel, so meanpool! falls through to 
+meanpool_direct!, which scalar-indexes and is therefore rejected on MtlArray.
+cuDNN provides the equivalent on the CUDA side, so this only shows up on Metal.
+
+    NNlib normalises 3-D and 4-D pooling to the 5-D method, so implementing 
+the 5-D case intercepts every call. Mean pooling is prod(kernel) strided reads
+and an accumulate: for a fixed window offset the source indices are strided and 
+non-overlapping, so this is race-free without atomics.
+
+    This is type piracy in the strict sense, as neither NNlib.meanpool! nor 
+MtlArray is ours, but it is the same shape as what NNlibCUDA does for CUDA, and 
+it lives in a Metal-triggered extension. =#
+
+
+using NNlib
+using NNlib: PoolDims, kernel_size, stride, dilation, padding, output_size
+
+function _mtlZeroPad(x::MtlArray{T,5}, p::NTuple{6,Int}) where {T}
+    all(iszero, p) && return x
+    s = size(x)
+    xp = similar(x, T, (s[1] + p[1] + p[2], s[2] + p[3] + p[4],
+                        s[3] + p[5] + p[6], s[4], s[5]))
+    fill!(xp, zero(T))
+    @views xp[p[1].+(1:s[1]), p[3].+(1:s[2]), p[5].+(1:s[3]), :, :] .= x
+    return xp
+end
+ 
+# indices of the window-offset `c` (0-based) along one spatial axis
+_poolIdx(c, st, dil, nOut) = (1 + c * dil):st:(1 + c * dil + st * (nOut - 1))
+ 
+function NNlib.meanpool!(y::MtlArray{T,5}, x::MtlArray{T,5}, pdims::PoolDims;
+                         alpha = true, beta = false, kwargs...) where {T}
+    kern = kernel_size(pdims)
+    st = stride(pdims)
+    dil = dilation(pdims)
+    o1, o2, o3 = output_size(pdims)
+    xp = _mtlZeroPad(x, padding(pdims))
+ 
+    acc = similar(y)
+    fill!(acc, zero(T))
+    for c3 = 0:kern[3]-1, c2 = 0:kern[2]-1, c1 = 0:kern[1]-1
+        i1 = _poolIdx(c1, st[1], dil[1], o1)
+        i2 = _poolIdx(c2, st[2], dil[2], o2)
+        i3 = _poolIdx(c3, st[3], dil[3], o3)
+        @views acc .+= xp[i1, i2, i3, :, :]
+    end
+    acc ./= T(prod(kern))
+ 
+    if iszero(beta)
+        y .= T(alpha) .* acc
+    else
+        y .= T(alpha) .* acc .+ T(beta) .* y
+    end
+    return y
+end
+ 
+function NNlib.∇meanpool!(dx::MtlArray{T,5}, dy::MtlArray{T,5}, y::MtlArray{T,5},
+                          x::MtlArray{T,5}, pdims::PoolDims;
+                          alpha = true, beta = false, kwargs...) where {T}
+    kern = kernel_size(pdims)
+    st = stride(pdims)
+    dil = dilation(pdims)
+    p = padding(pdims)
+    o1, o2, o3 = output_size(pdims)
+ 
+    scaled = dy ./ T(prod(kern))
+    padded = !all(iszero, p)
+    s = size(dx)
+    dxp = padded ?
+          fill!(similar(dx, T, (s[1] + p[1] + p[2], s[2] + p[3] + p[4],
+                                s[3] + p[5] + p[6], s[4], s[5])), zero(T)) :
+          fill!(dx, zero(T))
+ 
+    # each offset writes a strided, non-overlapping set of destinations, so the
+    # accumulate is safe; the overlap between offsets is serialised by the loop. 
+    for c3 = 0:kern[3]-1, c2 = 0:kern[2]-1, c1 = 0:kern[1]-1
+        i1 = _poolIdx(c1, st[1], dil[1], o1)
+        i2 = _poolIdx(c2, st[2], dil[2], o2)
+        i3 = _poolIdx(c3, st[3], dil[3], o3)
+        @views dxp[i1, i2, i3, :, :] .+= scaled
+    end
+ 
+    if padded
+        @views dx .= dxp[p[1].+(1:s[1]), p[3].+(1:s[2]), p[5].+(1:s[3]), :, :]
+    end
+    if !iszero(alpha) && alpha != true
+        dx .*= T(alpha)
+    end
+    return dx
+end
+
 end
\ No newline at end of file

From 7859327d8f4a56a59c48f586bf5be744abf736af Mon Sep 17 00:00:00 2001
From: JaredW40 <trumpetpro40@gmail.com>
Date: Wed, 19 Aug 2026 00:27:33 -0700
Subject: [PATCH 3/5] Updated testing to properly compare CPU and GPU. Removed
 BenchmarkTools. Updated files and tests accordingly.

---
 .github/workflows/CI.yml |  1 -
 Project.toml             |  4 +--
 ext/MetalExt.jl          |  1 -
 test/CUDATests.jl        | 26 ++++++-------------
 test/MetalPlanProbe.jl   | 56 ++++++++++++++++++++++++++++++++++++++++
 test/MetalTests.jl       | 23 +++++------------
 test/runtests.jl         | 23 ++++++++++++-----
 7 files changed, 88 insertions(+), 46 deletions(-)
 create mode 100644 test/MetalPlanProbe.jl

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 577ebca..3edfbe9 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -91,7 +91,6 @@ jobs:
             ])
             Pkg.add([
               PackageSpec(name = "Test"),
-              PackageSpec(name = "BenchmarkTools"),
               PackageSpec(name = "Wavelets"),
               PackageSpec(name = "Zygote"),
               PackageSpec(name = "Flux"),
diff --git a/Project.toml b/Project.toml
index 56a2c4b..8e26718 100644
--- a/Project.toml
+++ b/Project.toml
@@ -38,7 +38,6 @@ ScatteringPlotsExt = "Plots"
 [compat]
 AbstractFFTs = "1"
 Adapt = "4"
-BenchmarkTools = "1"
 ContinuousWavelets = "1"
 CUDA = "6.1"
 cuDNN = "6.1"
@@ -55,10 +54,9 @@ Wavelets = "0.9, 0.10"
 Zygote = "0.6, 0.7"
 
 [extras]
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "BenchmarkTools", "Logging", "Random"]
+test = ["Test", "Logging", "Random"]
diff --git a/ext/MetalExt.jl b/ext/MetalExt.jl
index 26efede..e9c484b 100644
--- a/ext/MetalExt.jl
+++ b/ext/MetalExt.jl
@@ -15,7 +15,6 @@ non-overlapping, so this is race-free without atomics.
 MtlArray is ours, but it is the same shape as what NNlibCUDA does for CUDA, and 
 it lives in a Metal-triggered extension. =#
 
-
 using NNlib
 using NNlib: PoolDims, kernel_size, stride, dilation, padding, output_size
 
diff --git a/test/CUDATests.jl b/test/CUDATests.jl
index 62283f4..7539b48 100644
--- a/test/CUDATests.jl
+++ b/test/CUDATests.jl
@@ -1,5 +1,3 @@
-using BenchmarkTools
-
 const cuda_available = @isdefined(CUDA) && CUDA.functional()
 
 @testset "GPU Tests (CUDA)" begin
@@ -150,7 +148,7 @@ const cuda_available = @isdefined(CUDA) && CUDA.functional()
 
         @testset "CPU/GPU timing" begin
             sizes = [256, 2048, 16384, 131072]
-            cpu_max_size = 16384
+            nSamples = 5 # Number of runs we test for CPU and GPU timing.
 
             for sz in sizes
                 GC.gc()
@@ -161,25 +159,17 @@ const cuda_available = @isdefined(CUDA) && CUDA.functional()
                 sstGPU = gpu(sst)
                 initGPU = gpu(init)
 
-                if sz > cpu_max_size
-                    CUDA.@sync sstGPU(initGPU)  # warmup
-                    GC.gc(); CUDA.reclaim()
-                    tGPU = @elapsed (CUDA.@sync sstGPU(initGPU))
-                else
-                    tGPU = @belapsed (CUDA.@sync $sstGPU($initGPU))
-                end
+                sst(init)
+                CUDA.@sync sstGPU(initGPU)
+                GC.gc(); CUDA.reclaim()
 
-                tCPU = @belapsed $sst($init)
-                speedup = tCPU / tGPU
-                @info "size=$sz" tCPU tGPU speedup
-                if sz >= 512
-                    @test tGPU < tCPU
-                end
+                tCPU = minimum(@elapsed(sst(init)) for _ = 1:nSamples)
+                tGPU = minimum(@elapsed(CUDA.@sync sstGPU(initGPU)) for _ = 1:nSamples)
+                @info "size=$sz" tCPU tGPU speedup = tCPU / tGPU
 
                 sstGPU = nothing
                 initGPU = nothing
-                GC.gc()
-                CUDA.reclaim()
+                GC.gc(); CUDA.reclaim()
             end
         end
     end
diff --git a/test/MetalPlanProbe.jl b/test/MetalPlanProbe.jl
new file mode 100644
index 0000000..2521d6e
--- /dev/null
+++ b/test/MetalPlanProbe.jl
@@ -0,0 +1,56 @@
+# test/MetalPlanProbe.jl
+#
+# Include this from MetalTests.jl BEFORE the model round-trip testset:
+#
+#     include("MetalPlanProbe.jl")
+#
+# Purpose: the stFlux round-trip test compares whole arrays, so when it fails it
+# tells you "these differ by 18%" and nothing about why. This isolates the plan
+# conversion itself. It runs in about a second and, whether it passes or fails,
+# prints everything needed to finish the fix without another guess.
+
+using Test, Metal, Adapt, AbstractFFTs, FFTW, FourierFilterFlux
+
+@testset "Metal FFT plan round-trip" begin
+    n, ch, batch = 16, 2, 3
+    xr = randn(Float32, n, ch, batch)
+    xc = complex(xr)
+
+    for (P, x, kind) in ((plan_rfft(xr, 1:1), xr, "r2c"),
+                         (plan_fft(xc, 1:1), xc, "c2c"))
+
+        G = Metal.mtl(P)
+
+        # --- what the Metal plan actually looks like -------------------------
+        # If the adapt patch guessed wrong about parameter order or property
+        # names, this block is the answer.
+        @info "plan $kind: cpu -> gpu" typeof(P) typeof(G)
+        @info "  sizes" cpuSize = size(P) gpuSize = size(G)
+        @info "  gpu properties" propertynames(G)
+        if hasproperty(G, :region)
+            @info "  regions" cpuRegion = P.region gpuRegion = G.region
+        else
+            @warn "  gpu plan has no :region property -- adapt will error; " *
+                  "use the correct name from propertynames above"
+        end
+
+        # Does size() on a Metal r2c plan report the real input length (16) or
+        # the complex output length (9)? The patch rebuilds with
+        # `zeros(Float32, size(p))`, which is only correct for the former.
+        if kind == "r2c"
+            @info "  r2c size check" expectedIfInputSize = n expectedIfOutputSize = n ÷ 2 + 1 actual = size(G, 1)
+        end
+
+        # --- the invariants that matter --------------------------------------
+        Q = adapt(Array, G)
+        @info "plan $kind: gpu -> cpu" typeof(Q) size(Q)
+
+        # kind is preserved: an rFFTWPlan must not come back as a cFFTWPlan
+        @test typeof(Q) == typeof(P)
+        # only the intended axis is transformed
+        @test Q.region == P.region
+        @test size(Q) == size(P)
+        # and the reconstructed plan actually computes the same thing
+        @test Q * x ≈ P * x
+    end
+end
diff --git a/test/MetalTests.jl b/test/MetalTests.jl
index c1acef3..95ab906 100644
--- a/test/MetalTests.jl
+++ b/test/MetalTests.jl
@@ -1,5 +1,3 @@
-using BenchmarkTools
-
 const metal_available = @isdefined(Metal) && Metal.functional()
 
 @testset "GPU Tests (Metal)" begin
@@ -150,7 +148,7 @@ const metal_available = @isdefined(Metal) && Metal.functional()
 
         @testset "CPU/GPU timing" begin
             sizes = [256, 2048, 16384, 131072]
-            cpu_max_size = 16384
+            nSamples = 5 # Number of runs we test for CPU and GPU timing. 
 
             for sz in sizes
                 GC.gc()
@@ -160,20 +158,13 @@ const metal_available = @isdefined(Metal) && Metal.functional()
                 sstGPU = gpu(sst)
                 initGPU = gpu(init)
 
-                if sz > cpu_max_size
-                    Metal.@sync sstGPU(initGPU)  # warmup
-                    GC.gc()
-                    tGPU = @elapsed (Metal.@sync sstGPU(initGPU))
-                else
-                    tGPU = @belapsed (Metal.@sync $sstGPU($initGPU))
-                end
+                sst(init)
+                Metal.@sync sstGPU(initGPU)
+                GC.gc()
 
-                tCPU = @belapsed $sst($init)
-                speedup = tCPU / tGPU
-                @info "size=$sz" tCPU tGPU speedup
-                if sz >= 512
-                    @test tGPU < tCPU
-                end
+                tCPU = minimum(@elapsed(sst(init)) for _ = 1:nSamples)
+                tGPU = minimum(@elapsed(Metal.@sync sstGPU(initGPU)) for _ = 1:nSamples)
+                @info "size=$sz" tCPU tGPU speedup = tCPU / tGPU
 
                 sstGPU = nothing
                 initGPU = nothing
diff --git a/test/runtests.jl b/test/runtests.jl
index 453abb8..46c60fb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -19,20 +19,29 @@ const GROUP = get(ENV, "GROUP", "All")
     include("2DTests.jl")
 
     if GROUP in ("All", "CUDA")
-        try
-            using CUDA, cuDNN, cuFFT
-            include("CUDATests.jl")
+        haveCUDA = try
+            @eval using CUDA, cuDNN, cuFFT
+            true
         catch e
             @info "CUDA/cuDNN/cuFFT not available in this environment -- skipping CUDATests.jl" exception=e
+            false
+        end
+        if haveCUDA
+            include("CUDATests.jl")
         end
     end
 
     if GROUP in ("All", "Metal")
-        try
-            using Metal
-            include("MetalTests.jl")
+        haveMetal = try
+            @eval using Metal
+            true
         catch e
-            @info "Metal not available in this environment -- skipping MetalTests.jl" exception=e
+            @info "Metal.jl not installed -- skipping Metal tests" exception = e
+            false
+        end
+        if haveMetal
+            include("MetalPlanProbe.jl")
+            include("MetalTests.jl")
         end
     end
 end
\ No newline at end of file

From 4a09246cb6e4ebaafe9709a965f796d95e127195 Mon Sep 17 00:00:00 2001
From: JaredW40 <trumpetpro40@gmail.com>
Date: Wed, 19 Aug 2026 01:45:27 -0700
Subject: [PATCH 4/5] Removed dedicated Metal probing tests and added smaller
 equivalents to both CUDATests.jl and MetalTests.jl.

---
 test/CUDATests.jl      | 11 +++++++++
 test/MetalPlanProbe.jl | 56 ------------------------------------------
 test/MetalTests.jl     | 11 +++++++++
 test/runtests.jl       |  3 +--
 4 files changed, 23 insertions(+), 58 deletions(-)
 delete mode 100644 test/MetalPlanProbe.jl

diff --git a/test/CUDATests.jl b/test/CUDATests.jl
index 7539b48..8970bf3 100644
--- a/test/CUDATests.jl
+++ b/test/CUDATests.jl
@@ -146,6 +146,17 @@ const cuda_available = @isdefined(CUDA) && CUDA.functional()
             end
         end
 
+        @testset "CUDA FFT plan round-trip" begin
+            x = randn(Float32, 16, 2, 3)
+            for (P, inp) in ((plan_rfft(x, 1:1), x), (plan_fft(complex(x), 1:1), complex(x)))
+                Q = adapt(Array, adapt(CuArray, P))
+                @test typeof(Q) == typeof(P)
+                @test Q.region == P.region
+                @test size(Q) == size(P)
+                @test Q * inp ≈ P * inp
+            end
+        end
+
         @testset "CPU/GPU timing" begin
             sizes = [256, 2048, 16384, 131072]
             nSamples = 5 # Number of runs we test for CPU and GPU timing.
diff --git a/test/MetalPlanProbe.jl b/test/MetalPlanProbe.jl
deleted file mode 100644
index 2521d6e..0000000
--- a/test/MetalPlanProbe.jl
+++ /dev/null
@@ -1,56 +0,0 @@
-# test/MetalPlanProbe.jl
-#
-# Include this from MetalTests.jl BEFORE the model round-trip testset:
-#
-#     include("MetalPlanProbe.jl")
-#
-# Purpose: the stFlux round-trip test compares whole arrays, so when it fails it
-# tells you "these differ by 18%" and nothing about why. This isolates the plan
-# conversion itself. It runs in about a second and, whether it passes or fails,
-# prints everything needed to finish the fix without another guess.
-
-using Test, Metal, Adapt, AbstractFFTs, FFTW, FourierFilterFlux
-
-@testset "Metal FFT plan round-trip" begin
-    n, ch, batch = 16, 2, 3
-    xr = randn(Float32, n, ch, batch)
-    xc = complex(xr)
-
-    for (P, x, kind) in ((plan_rfft(xr, 1:1), xr, "r2c"),
-                         (plan_fft(xc, 1:1), xc, "c2c"))
-
-        G = Metal.mtl(P)
-
-        # --- what the Metal plan actually looks like -------------------------
-        # If the adapt patch guessed wrong about parameter order or property
-        # names, this block is the answer.
-        @info "plan $kind: cpu -> gpu" typeof(P) typeof(G)
-        @info "  sizes" cpuSize = size(P) gpuSize = size(G)
-        @info "  gpu properties" propertynames(G)
-        if hasproperty(G, :region)
-            @info "  regions" cpuRegion = P.region gpuRegion = G.region
-        else
-            @warn "  gpu plan has no :region property -- adapt will error; " *
-                  "use the correct name from propertynames above"
-        end
-
-        # Does size() on a Metal r2c plan report the real input length (16) or
-        # the complex output length (9)? The patch rebuilds with
-        # `zeros(Float32, size(p))`, which is only correct for the former.
-        if kind == "r2c"
-            @info "  r2c size check" expectedIfInputSize = n expectedIfOutputSize = n ÷ 2 + 1 actual = size(G, 1)
-        end
-
-        # --- the invariants that matter --------------------------------------
-        Q = adapt(Array, G)
-        @info "plan $kind: gpu -> cpu" typeof(Q) size(Q)
-
-        # kind is preserved: an rFFTWPlan must not come back as a cFFTWPlan
-        @test typeof(Q) == typeof(P)
-        # only the intended axis is transformed
-        @test Q.region == P.region
-        @test size(Q) == size(P)
-        # and the reconstructed plan actually computes the same thing
-        @test Q * x ≈ P * x
-    end
-end
diff --git a/test/MetalTests.jl b/test/MetalTests.jl
index 95ab906..5d78943 100644
--- a/test/MetalTests.jl
+++ b/test/MetalTests.jl
@@ -146,6 +146,17 @@ const metal_available = @isdefined(Metal) && Metal.functional()
             end
         end
 
+        @testset "Metal FFT plan round-trip" begin
+            x = randn(Float32, 16, 2, 3)
+            for (P, inp) in ((plan_rfft(x, 1:1), x), (plan_fft(complex(x), 1:1), complex(x)))
+                Q = adapt(Array, Metal.mtl(P))
+                @test typeof(Q) == typeof(P)
+                @test Q.region == P.region
+                @test size(Q) == size(P)
+                @test Q * inp ≈ P * inp
+            end
+        end
+
         @testset "CPU/GPU timing" begin
             sizes = [256, 2048, 16384, 131072]
             nSamples = 5 # Number of runs we test for CPU and GPU timing. 
diff --git a/test/runtests.jl b/test/runtests.jl
index 46c60fb..a942560 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2,7 +2,7 @@
 using ScatteringTransform
 using ContinuousWavelets
 using AbstractFFTs, FFTW
-using Test, LinearAlgebra, Statistics
+using Test, LinearAlgebra, Statistics, Adapt
 using Flux, FourierFilterFlux, MonogenicFilterFlux
 using Zygote
 
@@ -40,7 +40,6 @@ const GROUP = get(ENV, "GROUP", "All")
             false
         end
         if haveMetal
-            include("MetalPlanProbe.jl")
             include("MetalTests.jl")
         end
     end

From 068ec014d5aed831fcfd78bcecfb99978c29673a Mon Sep 17 00:00:00 2001
From: JaredW40 <trumpetpro40@gmail.com>
Date: Wed, 19 Aug 2026 13:45:43 -0700
Subject: [PATCH 5/5] Removed dependency on cuDNN for CUDA.

---
 .github/workflows/CI.yml |  2 --
 Project.toml             |  6 +-----
 ext/CUDAExt.jl           | 14 ++++++--------
 test/runtests.jl         |  4 ++--
 4 files changed, 9 insertions(+), 17 deletions(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 3edfbe9..2320147 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -111,9 +111,7 @@ jobs:
             using Pkg
             Pkg.add([
               PackageSpec(name = "CUDA"),
-              PackageSpec(name = "cuDNN"),
               PackageSpec(name = "cuFFT"),
-              PackageSpec(name = "CUDNN_jll", version = "9.20"),
             ])'
       - name: Add Metal (macOS only)
         if: runner.os == 'macOS'
diff --git a/Project.toml b/Project.toml
index 8e26718..deb57c1 100644
--- a/Project.toml
+++ b/Project.toml
@@ -24,14 +24,12 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
-CUDNN_jll = "62b44479-cb7b-5706-934f-f13b2eb2e645"
-cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 cuFFT = "533571aa-0936-420e-b4be-9c66f5f626ca"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [extensions]
-CUDAExt = ["CUDA", "cuDNN", "cuFFT", "CUDNN_jll"]
+CUDAExt = ["CUDA", "cuFFT"]
 MetalExt = "Metal"
 ScatteringPlotsExt = "Plots"
 
@@ -40,8 +38,6 @@ AbstractFFTs = "1"
 Adapt = "4"
 ContinuousWavelets = "1"
 CUDA = "6.1"
-cuDNN = "6.1"
-CUDNN_jll = "~9.20"
 cuFFT = "6.1"
 FFTW = "1"
 Flux = "0.13, 0.14, 0.15, 0.16"
diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl
index c930723..2cca59c 100644
--- a/ext/CUDAExt.jl
+++ b/ext/CUDAExt.jl
@@ -8,13 +8,11 @@ module CUDAExt
     FourierFilterFlux is loaded.
 
     This file's real job is triggering activation: `[extensions] CUDAExt =
-    ["CUDA", "cuDNN", "cuFFT"]` in Project.toml means all three must be
-    loaded before this extension (and therefore this package's CUDA support)
-    activates at all, so this `using` line is what actually makes that
-    happen, not a formality. It also guarantees FourierFilterFlux's and
-    ContinuousWavelets' own CUDA extensions are active by the time anything
-    in ScatteringTransform runs on the GPU, since their trigger sets
-    (CUDA+cuFFT, and CUDA+cuDNN+cuFFT, respectively) are subsets of this one. =#
-using ScatteringTransform, CUDA, cuDNN, cuFFT
+    ["CUDA", "cuFFT"]` in Project.toml means both must be loaded before this 
+    extension (and therefore this package's CUDA support) activates at all. 
+    It also guarantees FourierFilterFlux's and ContinuousWavelets' own CUDA 
+    extensions are active by the time anything in ScatteringTransform runs on 
+    the GPU, since their trigger sets are the same as this one. =#
+using ScatteringTransform, CUDA, cuFFT
 
 end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index a942560..0587261 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -20,10 +20,10 @@ const GROUP = get(ENV, "GROUP", "All")
 
     if GROUP in ("All", "CUDA")
         haveCUDA = try
-            @eval using CUDA, cuDNN, cuFFT
+            @eval using CUDA, cuFFT
             true
         catch e
-            @info "CUDA/cuDNN/cuFFT not available in this environment -- skipping CUDATests.jl" exception=e
+            @info "CUDA/cuFFT not available in this environment -- skipping CUDATests.jl" exception=e
             false
         end
         if haveCUDA