Remove old example for now

Author: avik-pal

.buildkite/pipeline.yml

@@ -7,9 +7,6 @@ steps:
           test_args: "--quickfail"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-          dirs:
-            - src
-            - ext
     agents:
       queue: "juliagpu"
       cuda: "*"
@@ -30,9 +27,6 @@ steps:
           test_args: "--quickfail"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-          dirs:
-            - src
-            - ext
     env:
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"

.github/workflows/CI.yml

@@ -41,8 +41,6 @@ jobs:
           RETESTITEMS_NWORKERS: 4
           RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
-        with:
-          directories: src,ext
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info

LocalPreferences.toml

@@ -0,0 +1,2 @@
+[LuxTestUtils]
+target_modules = ["LuxNeuralOperators", "Lux", "LuxLib"]

codecov.yml

@@ -0,0 +1,3 @@
+codecov:
+  notify:
+    wait_for_ci: false

examples/Burgers/main.jl

@@ -28,7 +28,7 @@ end
 
 @inline truncate_modes(ft::FourierTransform, x_fft::AbstractArray) = low_pass(ft, x_fft)
 
-function inverse(ft::FourierTransform, x_fft::AbstractArray{T, N},
-        M::NTuple{N, Int64}) where {T, N}
+function inverse(
+        ft::FourierTransform, x_fft::AbstractArray{T, N}, M::NTuple{N, Int64}) where {T, N}
     return real(irfft(x_fft, first(M), 1:ndims(ft)))
 end

examples/Project.toml

@@ -1,20 +1,23 @@
 @testitem "Fourier Neural Operator" setup=[SharedTestSetup] begin
     @testset "BACKEND: $(mode)" for (mode, aType, dev, ongpu) in MODES
-        rng = get_default_rng(mode)
+        rng = get_stable_rng()
 
         setups = [
-            (modes=(16,), chs=(2, 64, 64, 64, 64, 64, 128, 1), x_size=(2, 1024, 5),
-                y_size=(1, 1024, 5), permuted=Val(false)),
-            (modes=(16,), chs=(2, 64, 64, 64, 64, 64, 128, 1), x_size=(1024, 2, 5),
-                y_size=(1024, 1, 5), permuted=Val(true))]
+            (modes=(16,), chs=(2, 64, 64, 64, 64, 64, 128, 1),
+                x_size=(2, 1024, 5), y_size=(1, 1024, 5), permuted=Val(false)),
+            (modes=(16,), chs=(2, 64, 64, 64, 64, 64, 128, 1),
+                x_size=(1024, 2, 5), y_size=(1024, 1, 5), permuted=Val(true))]
 
         @testset "$(length(setup.modes))D: permuted = $(setup.permuted)" for setup in setups
             fno = FourierNeuralOperator(; setup.chs, setup.modes, setup.permuted)
 
-            x = rand(rng, Float32, setup.x_size...)
-            y = rand(rng, Float32, setup.y_size...)
+            x = rand(rng, Float32, setup.x_size...) |> aType
+            y = rand(rng, Float32, setup.y_size...) |> aType
 
-            ps, st = Lux.setup(rng, fno)
+            ps, st = Lux.setup(rng, fno) |> dev
+
+            @inferred fno(x, ps, st)
+            @jet fno(x, ps, st)
 
             @test size(first(fno(x, ps, st))) == setup.y_size
 

examples/common.jl

@@ -1,6 +1,6 @@
 @testitem "SpectralConv & SpectralKernel" setup=[SharedTestSetup] begin
     @testset "BACKEND: $(mode)" for (mode, aType, dev, ongpu) in MODES
-        rng = get_default_rng(mode)
+        rng = get_stable_rng()
 
         opconv = [SpectralConv, SpectralKernel]
         setups = [
@@ -11,7 +11,7 @@
             (; m=(10, 10), permuted=Val(true),
                 x_size=(22, 22, 1, 5), y_size=(22, 22, 64, 5))]
 
-        @testset "$(op) $(length(setup.modes))D: permuted = $(setup.permuted)" for setup in setups,
+        @testset "$(op) $(length(setup.m))D: permuted = $(setup.permuted)" for setup in setups,
             op in opconv
 
             p = Lux.__unwrap_val(setup.permuted)
@@ -22,13 +22,14 @@
             l1 = p ? Conv(ntuple(_ -> 1, length(setup.m)), in_chs => first(ch)) :
                  Dense(in_chs => first(ch))
             m = Chain(l1, op(ch, setup.m; setup.permuted))
-            ps, st = Lux.setup(rng, m)
+            ps, st = Lux.setup(rng, m) |> dev
 
-            x = rand(rng, Float32, setup.x_size...)
+            x = rand(rng, Float32, setup.x_size...) |> aType
             @test size(first(m(x, ps, st))) == setup.y_size
             @inferred m(x, ps, st)
+            @jet m(x, ps, st)
 
-            data = [(x, rand(rng, Float32, setup.y_size...))]
+            data = [(x, aType(rand(rng, Float32, setup.y_size...)))]
             l2, l1 = train!(m, ps, st, data; epochs=10)
             @test l2 < l1
         end

src/transform.jl

@@ -43,7 +43,7 @@ default_loss_function(model, ps, x, y) = mean(abs2, y .- model(x, ps))
 train!(args...; kwargs...) = train!(default_loss_function, args...; kwargs...)
 
 function train!(loss, model, ps, st, data; epochs=10)
-    m = StatefulLuxLayer(model, ps, st)
+    m = StatefulLuxLayer{true}(model, ps, st)
 
     l1 = loss(m, ps, first(data)...)
     st_opt = Optimisers.setup(Adam(0.01f0), ps)