feat: update to newest versions

Author: avik-pal

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,12 @@ 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"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
 MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+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,10 +22,12 @@ Documenter = "1.5.0"
 ExplicitImports = "1.9.0"
 Hwloc = "3.2.0"
 InteractiveUtils = "<0.0.1, 1"
+Lux = "0.5.62"
 LuxTestUtils = "1.1.2"
 MLDataDevices = "1.0.0"
 Optimisers = "0.3.3"
 Pkg = "1.10"
+Random = "1.10"
 Reexport = "1.2.2"
 ReTestItems = "1.24.0"
 StableRNGs = "1.0.2"

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/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))