feat: update to newest versions

Author: avik-pal

.buildkite/testing.yml

@@ -7,8 +7,6 @@ steps:
           test_args: "--quickfail"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-    commands: |
-      printf "[LuxTestUtils]\ntarget_modules = [\"NeuralOperators\", \"Lux\", \"LuxLib\"]\n[LuxLib]\ninstability_check = \"error\"\n[LuxCore]\ninstability_check = \"error\"\n" > LocalPreferences.toml
     agents:
       queue: "juliagpu"
       cuda: "*"
@@ -29,8 +27,6 @@ steps:
           test_args: "--quickfail"
       - JuliaCI/julia-coverage#v1:
           codecov: true
-    commands: |
-      printf "[LuxTestUtils]\ntarget_modules = [\"NeuralOperators\", \"Lux\", \"LuxLib\"]\n[LuxLib]\ninstability_check = \"error\"\n[LuxCore]\ninstability_check = \"error\"\n" > LocalPreferences.toml
     env:
       JULIA_AMDGPU_CORE_MUST_LOAD: "1"
       JULIA_AMDGPU_HIP_MUST_LOAD: "1"

.github/workflows/CI.yml

@@ -35,16 +35,6 @@ jobs:
           - windows-latest
     steps:
       - uses: actions/checkout@v4
-      - uses: DamianReeves/write-file-action@master
-        with:
-          path: "LocalPreferences.toml"
-          contents: |
-            [LuxTestUtils]
-            target_modules = ["NeuralOperators", "Lux", "LuxLib"]
-            [LuxLib]
-            instability_check = "error"
-            [LuxCore]
-            instability_check = "error"
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
@@ -78,16 +68,6 @@ jobs:
         version: ["1"]
     steps:
       - uses: actions/checkout@v4
-      - uses: DamianReeves/write-file-action@master
-        with:
-          path: "LocalPreferences.toml"
-          contents: |
-            [LuxTestUtils]
-            target_modules = ["NeuralOperators", "Lux", "LuxLib"]
-            [LuxLib]
-            instability_check = "error"
-            [LuxCore]
-            instability_check = "error"
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}

Project.toml

@@ -22,9 +22,9 @@ ChainRulesCore = "1.24.0"
 ConcreteStructs = "0.2.3"
 FFTW = "1.8.0"
 Lux = "0.5.62"
-LuxCore = "0.1.15"
+LuxCore = "0.1.21"
 LuxLib = "0.3.40"
-NNlib = "0.9.17"
+NNlib = "0.9.21"
 Random = "1.10"
 Reexport = "1.2.2"
 WeightInitializers = "1"

test/Project.toml

@@ -4,10 +4,15 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -20,12 +25,17 @@ Documenter = "1.5.0"
 ExplicitImports = "1.9.0"
 Hwloc = "3.2.0"
 InteractiveUtils = "<0.0.1, 1"
+Lux = "0.5.62"
+LuxCore = "0.1.21"
+LuxLib = "0.3.40"
 LuxTestUtils = "1.1.2"
 MLDataDevices = "1.0.0"
 Optimisers = "0.3.3"
 Pkg = "1.10"
-Reexport = "1.2.2"
+Preferences = "1"
+Random = "1.10"
 ReTestItems = "1.24.0"
+Reexport = "1.2.2"
 StableRNGs = "1.0.2"
 Test = "1.10"
 Zygote = "0.6.70"

test/deeponet_tests.jl

@@ -50,7 +50,8 @@
             @test setup.out_size == size(pred)
 
             __f = (u, y, ps) -> sum(abs2, first(deeponet((u, y), ps, st)))
-            test_gradients(__f, u, y, ps; atol=1.0f-3, rtol=1.0f-3)
+            test_gradients(
+                __f, u, y, ps; atol=1.0f-3, rtol=1.0f-3, skip_backends=[AutoEnzyme()])
         end
 
         @testset "Embedding layer mismatch" begin
@@ -62,9 +63,6 @@
 
             ps, st = Lux.setup(rng, deeponet) |> dev
             @test_throws ArgumentError deeponet((u, y), ps, st)
-
-            __f = (u, y, ps) -> sum(abs2, first(deeponet((u, y), ps, st)))
-            test_gradients(__f, u, y, ps; atol=1.0f-3, rtol=1.0f-3)
         end
     end
 end

test/fno_tests.jl

@@ -29,7 +29,8 @@
             end broken=broken
 
             __f = (x, ps) -> sum(abs2, first(fno(x, ps, st)))
-            test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3)
+            test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3,
+                skip_backends=[AutoEnzyme(), AutoTracker(), AutoReverseDiff()])
         end
     end
 end

test/layers_tests.jl

@@ -38,7 +38,8 @@
             end broken=broken
 
             __f = (x, ps) -> sum(abs2, first(m(x, ps, st)))
-            test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3)
+            test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3,
+                skip_backends=[AutoEnzyme(), AutoTracker(), AutoReverseDiff()])
         end
     end
 end

test/qa_tests.jl

@@ -9,7 +9,8 @@ end
 @testitem "Aqua: Quality Assurance" tags=[:qa] begin
     using Aqua
 
-    Aqua.test_all(NeuralOperators)
+    Aqua.test_all(NeuralOperators; ambiguities=false)
+    Aqua.test_ambiguities(NeuralOperators; recursive=false)
 end
 
 @testitem "Explicit Imports: Quality Assurance" tags=[:qa] begin

test/runtests.jl

@@ -1,6 +1,9 @@
-using ReTestItems, Pkg, Test
-using InteractiveUtils, Hwloc
-using NeuralOperators
+using Preferences
+
+Preferences.set_preferences!("LuxLib", "instability_check" => "error")
+Preferences.set_preferences!("LuxCore", "instability_check" => "error")
+
+using ReTestItems, Pkg, Test, InteractiveUtils, Hwloc, NeuralOperators
 
 const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "all"))
 

test/shared_testsetup.jl

@@ -4,6 +4,8 @@ import Reexport: @reexport
 @reexport using Lux, Zygote, Optimisers, Random, StableRNGs, LuxTestUtils
 using MLDataDevices
 
+LuxTestUtils.jet_target_modules!(["NeuralOperators", "Lux", "LuxLib"])
+
 const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "All"))
 
 if BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda"
@@ -37,9 +39,9 @@ train!(args...; kwargs...) = train!(MSELoss(), AutoZygote(), args...; kwargs...)
 function train!(loss, backend, model, ps, st, data; epochs=10)
     l1 = loss(model, ps, st, first(data))
 
-    tstate = Lux.Experimental.TrainState(model, ps, st, Adam(0.01f0))
+    tstate = Training.TrainState(model, ps, st, Adam(0.01f0))
     for _ in 1:epochs, (x, y) in data
-        _, _, _, tstate = Lux.Experimental.single_train_step!(backend, loss, (x, y), tstate)
+        _, _, _, tstate = Training.single_train_step!(backend, loss, (x, y), tstate)
     end
 
     l2 = loss(model, ps, st, first(data))