feat: update to newest versions

Author: avik-pal

Project.toml

@@ -10,47 +10,22 @@ ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
-LuxDeviceUtils = "34f89e08-e1d5-43b4-8944-0b49ac560553"
+LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 WeightInitializers = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 
 [compat]
-Aqua = "0.8.7"
 ArgCheck = "2.3.0"
 ChainRulesCore = "1.24.0"
 ConcreteStructs = "0.2.3"
-Documenter = "1.4.1"
-ExplicitImports = "1.9.0"
 FFTW = "1.8.0"
-Lux = "0.5.56"
+Lux = "0.5.62"
 LuxCore = "0.1.15"
-LuxDeviceUtils = "0.1.24"
-LuxTestUtils = "0.1.15"
+LuxLib = "0.3.40"
 NNlib = "0.9.17"
-Optimisers = "0.3.3"
-Pkg = "1.10"
 Random = "1.10"
-ReTestItems = "1.24.0"
 Reexport = "1.2.2"
-StableRNGs = "1.0.2"
-Test = "1.10"
-WeightInitializers = "0.1.7, 1"
-Zygote = "0.6.70"
+WeightInitializers = "1"
 julia = "1.10"
-
-[extras]
-Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
-Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
-ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
-LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
-Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
-Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
-ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
-StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
-
-[targets]
-test = ["Aqua", "Documenter", "ExplicitImports", "LuxTestUtils", "Optimisers", "Pkg", "ReTestItems", "StableRNGs", "Test", "Zygote"]

src/NeuralOperators.jl

@@ -6,8 +6,8 @@ using ConcreteStructs: @concrete
 using FFTW: FFTW, irfft, rfft
 using Lux
 using LuxCore: LuxCore, AbstractExplicitLayer
-using LuxDeviceUtils: get_device, LuxAMDGPUDevice
-using NNlib: NNlib, ⊠, batched_adjoint
+using LuxLib: batched_matmul
+using NNlib: NNlib, batched_adjoint
 using Random: Random, AbstractRNG
 using Reexport: @reexport
 

src/functional.jl

@@ -12,8 +12,8 @@ end
     x_size = size(x_tr)
     x_flat = reshape(x_tr, :, x_size[N - 1], x_size[N])
 
-    x_flat_t = permutedims(x_flat, (2, 3, 1))                              # i x b x m
-    x_weighted = permutedims(__batched_mul(weights, x_flat_t), (3, 1, 2))  # m x o x b
+    x_flat_t = permutedims(x_flat, (2, 3, 1))                               # i x b x m
+    x_weighted = permutedims(batched_matmul(weights, x_flat_t), (3, 1, 2))  # m x o x b
 
     return reshape(x_weighted, x_size[1:(N - 2)]..., size(x_weighted)[2:3]...)
 end

src/utils.jl

@@ -1,15 +1,3 @@
-# Temporarily capture certain calls like AMDGPU for ComplexFloats
-@inline __batched_mul(x, y) = __batched_mul(x, y, get_device((x, y)))
-@inline function __batched_mul(
-        x::AbstractArray{<:Number, 3}, y::AbstractArray{<:Number, 3}, _)
-    return x ⊠ y
-end
-@inline function __batched_mul(
-        x::AbstractArray{<:Complex, 3}, y::AbstractArray{<:Complex, 3}, ::LuxAMDGPUDevice)
-    # FIXME: This is not good for performance but that is okay for now
-    return stack(*, eachslice(x; dims=3), eachslice(y; dims=3))
-end
-
 @inline function __project(b::AbstractArray{T1, 2}, t::AbstractArray{T2, 3},
         additional::Nothing) where {T1, T2}
     # b : p x nb
@@ -25,7 +13,7 @@ end
     if size(b, 2) == 1 || size(t, 2) == 1
         return sum(b .* t; dims=1) # 1 x N x nb
     else
-        return __batched_mul(batched_adjoint(b), t) # u x N x b
+        return batched_matmul(batched_adjoint(b), t) # u x N x b
     end
 end
 

test/Project.toml

@@ -0,0 +1,31 @@
+[deps]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+Aqua = "0.8.7"
+Documenter = "1.5.0"
+ExplicitImports = "1.9.0"
+Hwloc = "3.2.0"
+InteractiveUtils = "<0.0.1, 1"
+LuxTestUtils = "1.1.2"
+MLDataDevices = "1.0.0"
+Optimisers = "0.3.3"
+Pkg = "1.10"
+Reexport = "1.2.2"
+ReTestItems = "1.24.0"
+StableRNGs = "1.0.2"
+Test = "1.10"
+Zygote = "0.6.70"

test/deeponet_tests.jl

@@ -48,6 +48,9 @@
 
             pred = first(deeponet((u, y), ps, st))
             @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=1f-3, rtol=1f-3)
         end
 
         @testset "Embedding layer mismatch" begin
@@ -59,6 +62,9 @@
 
             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=1f-3, rtol=1f-3)
         end
     end
 end

test/fno_tests.jl

@@ -27,6 +27,9 @@
                 l2, l1 = train!(fno, ps, st, data; epochs=10)
                 l2 < l1
             end broken=broken
+
+            __f = (x, ps) -> sum(abs2, first(fno(x, ps, st)))
+            test_gradients(__f, x, ps; atol=1f-3, rtol=1f-3)
         end
     end
 end

test/layers_tests.jl

@@ -36,6 +36,9 @@
                 l2, l1 = train!(m, ps, st, data; epochs=10)
                 l2 < l1
             end broken=broken
+
+            __f = (x, ps) -> sum(abs2, first(m(x, ps, st)))
+            test_gradients(__f, x, ps; atol=1f-3, rtol=1f-3)
         end
     end
 end

test/runtests.jl

@@ -1,4 +1,6 @@
-using ReTestItems, Pkg, ReTestItems, Test
+using ReTestItems, Pkg, Test
+using InteractiveUtils, Hwloc
+using NeuralOperators
 
 const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "all"))
 
@@ -14,6 +16,13 @@ if !isempty(EXTRA_PKGS)
     Pkg.instantiate()
 end
 
+const RETESTITEMS_NWORKERS = parse(
+    Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 16))))
+const RETESTITEMS_NWORKER_THREADS = parse(Int,
+    get(ENV, "RETESTITEMS_NWORKER_THREADS",
+        string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1))))
+
 @testset "NeuralOperators.jl Tests" begin
-    ReTestItems.runtests(@__DIR__)
+    ReTestItems.runtests(NeuralOperators; nworkers=RETESTITEMS_NWORKERS,
+        nworker_threads=RETESTITEMS_NWORKER_THREADS, testitem_timeout=3600)
 end

test/shared_testsetup.jl

@@ -1,8 +1,8 @@
 @testsetup module SharedTestSetup
 import Reexport: @reexport
 
-@reexport using Lux, Zygote, Optimisers, Random, StableRNGs
-using LuxTestUtils: @jet, @test_gradients
+@reexport using Lux, Zygote, Optimisers, Random, StableRNGs, LuxTestUtils
+using MLDataDevices
 
 const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "All"))
 
@@ -17,18 +17,18 @@ end
 cpu_testing() = BACKEND_GROUP == "all" || BACKEND_GROUP == "cpu"
 function cuda_testing()
     return (BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda") &&
-           LuxDeviceUtils.functional(LuxCUDADevice)
+           MLDataDevices.functional(CUDADevice)
 end
 function amdgpu_testing()
     return (BACKEND_GROUP == "all" || BACKEND_GROUP == "amdgpu") &&
-           LuxDeviceUtils.functional(LuxAMDGPUDevice)
+           MLDataDevices.functional(AMDGPUDevice)
 end
 
 const MODES = begin
     modes = []
-    cpu_testing() && push!(modes, ("CPU", Array, LuxCPUDevice(), false))
-    cuda_testing() && push!(modes, ("CUDA", CuArray, LuxCUDADevice(), true))
-    amdgpu_testing() && push!(modes, ("AMDGPU", ROCArray, LuxAMDGPUDevice(), true))
+    cpu_testing() && push!(modes, ("CPU", Array, CPUDevice(), false))
+    cuda_testing() && push!(modes, ("CUDA", CuArray, CUDADevice(), true))
+    amdgpu_testing() && push!(modes, ("AMDGPU", ROCArray, AMDGPUDevice(), true))
     modes
 end
 
@@ -47,7 +47,7 @@ function train!(loss, backend, model, ps, st, data; epochs=10)
     return l2, l1
 end
 
-export @jet, @test_gradients, check_approx
+export check_approx
 export BACKEND_GROUP, MODES, cpu_testing, cuda_testing, amdgpu_testing, train!
 
 end