test: reduce lux groups

.JuliaFormatter.toml

@@ -1,4 +1,4 @@
-style = "blue"
+style = "sciml"
 whitespace_in_kwargs = false
 margin = 92
 indent = 4

.github/workflows/CI.yml

@@ -39,15 +39,10 @@ jobs:
           - windows-latest
         test_group:
           - "core_layers"
-          - "contrib"
-          - "helpers"
-          - "distributed"
           - "normalize_layers"
-          - "others"
-          - "autodiff"
           - "recurrent_layers"
-          - "eltype_match"
-          - "fluxcompat"
+          - "autodiff"
+          - "misc"
           - "reactant"
     steps:
       - uses: actions/checkout@v4

.github/workflows/CIPreRelease.yml

@@ -37,15 +37,10 @@ jobs:
           - ubuntu-latest
         test_group:
           - "core_layers"
-          - "contrib"
-          - "helpers"
-          - "distributed"
           - "normalize_layers"
-          - "others"
-          - "autodiff"
           - "recurrent_layers"
-          - "eltype_match"
-          - "fluxcompat"
+          - "autodiff"
+          - "misc"
           - "reactant"
     steps:
       - uses: actions/checkout@v4
@@ -62,8 +57,30 @@ jobs:
             ${{ runner.os }}-test-${{ env.cache-name }}-
             ${{ runner.os }}-test-
             ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
+      - name: "Install Dependencies"
+        run: |
+          import Pkg
+          dev_pkgs = Pkg.PackageSpec[]
+          for pkg in ("lib/LuxCore", "lib/MLDataDevices", "lib/WeightInitializers", "lib/LuxLib",)
+            push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
+          end
+          Pkg.develop(dev_pkgs)
+          Pkg.Registry.update()
+          Pkg.instantiate()
+          Pkg.activate("test")
+          dev_pkgs = Pkg.PackageSpec[]
+          for pkg in ("lib/LuxTestUtils", "lib/LuxLib", "lib/MLDataDevices", "lib/LuxCore", ".")
+            push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
+          end
+          Pkg.develop(dev_pkgs)
+          Pkg.instantiate()
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
+      - name: "Run Tests"
+        run: |
+          import Pkg, Lux
+          dir = dirname(pathof(Lux))
+          include(joinpath(dir, "../test/runtests.jl"))
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=test {0}
         env:
           LUX_TEST_GROUP: ${{ matrix.test_group }}
           BACKEND_GROUP: "CPU"

.github/workflows/CI_LuxLib.yml

@@ -132,16 +132,6 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
-      matrix:
-        test_group:
-          - "conv"
-          - "dense"
-          - "normalization"
-          - "misc"
-        blas_backend:
-          - "default"
-        loopvec:
-          - "true"
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -174,9 +164,9 @@ jobs:
           include(joinpath(dir, "../test/runtests.jl"))
         shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/LuxLib/test {0}
         env:
-          LUXLIB_TEST_GROUP: ${{ matrix.test_group }}
-          LUXLIB_BLAS_BACKEND: ${{ matrix.blas_backend }}
-          LUXLIB_LOAD_LOOPVEC: ${{ matrix.loopvec }}
+          LUXLIB_TEST_GROUP: "all"
+          LUXLIB_BLAS_BACKEND: "default"
+          LUXLIB_LOAD_LOOPVEC: "true"
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: lib/LuxLib/src,lib/LuxLib/ext,lib/LuxCore/src,lib/LuxCore/ext,lib/MLDataDevices/src,lib/MLDataDevices/ext,lib/LuxTestUtils/src

Project.toml

@@ -79,7 +79,7 @@ DispatchDoctor = "0.4.12"
 Enzyme = "0.13.1"
 EnzymeCore = "0.8.1"
 FastClosures = "0.3.2"
-Flux = "0.14.20"
+Flux = "0.14.25"
 ForwardDiff = "0.10.36"
 FunctionWrappers = "1.1.3"
 Functors = "0.4.12"

test/contrib/debug_tests.jl

@@ -1,4 +1,4 @@
-@testitem "Debugging Tools: DimensionMismatch" setup=[SharedTestSetup] tags=[:contrib] begin
+@testitem "Debugging Tools: DimensionMismatch" setup=[SharedTestSetup] tags=[:misc] begin
     using Logging
 
     rng = StableRNG(12345)
@@ -43,7 +43,7 @@
     end
 end
 
-@testitem "Debugging Tools: NaN" setup=[SharedTestSetup] tags=[:contrib] begin
+@testitem "Debugging Tools: NaN" setup=[SharedTestSetup] tags=[:misc] begin
     using Logging, ChainRulesCore
     import ChainRulesCore as CRC
 

test/contrib/freeze_tests.jl

@@ -1,4 +1,4 @@
-@testitem "All Parameter Freezing" setup=[SharedTestSetup] tags=[:contrib] begin
+@testitem "All Parameter Freezing" setup=[SharedTestSetup] tags=[:misc] begin
     rng = StableRNG(12345)
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES
@@ -63,7 +63,7 @@
     end
 end
 
-@testitem "Partial Freezing" setup=[SharedTestSetup] tags=[:contrib] begin
+@testitem "Partial Freezing" setup=[SharedTestSetup] tags=[:misc] begin
     using Lux.Experimental: FrozenLayer
 
     rng = StableRNG(12345)

test/contrib/map_tests.jl

@@ -1,4 +1,4 @@
-@testitem "Layer Map" setup=[SharedTestSetup] tags=[:contrib] begin
+@testitem "Layer Map" setup=[SharedTestSetup] tags=[:misc] begin
     using Setfield, Functors
 
     function occurs_in(kp::KeyPath, x::KeyPath)

test/contrib/share_parameters_tests.jl

@@ -1,4 +1,4 @@
-@testitem "Parameter Sharing" setup=[SharedTestSetup] tags=[:contrib] begin
+@testitem "Parameter Sharing" setup=[SharedTestSetup] tags=[:misc] begin
     rng = StableRNG(12345)
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES

test/enzyme_tests.jl

@@ -41,7 +41,8 @@ const MODELS_LIST  = [
     (Chain(Conv((4, 4), 2 => 2; pad=SamePad()), MeanPool((5, 5); pad=SamePad())), rand(Float32, 5, 5, 2, 2)),
     (Chain(Conv((3, 3), 2 => 3, relu; pad=SamePad()), MaxPool((2, 2))), rand(Float32, 5, 5, 2, 2)),
     (Maxout(() -> Dense(5 => 4, tanh), 3), randn(Float32, 5, 2)),
-    (Bilinear((2, 2) => 3), randn(Float32, 2, 3)),
+    # XXX: https://github.com/LuxDL/Lux.jl/issues/1024
+    # (Bilinear((2, 2) => 3), randn(Float32, 2, 3)),
     (SkipConnection(Dense(2 => 2), vcat), randn(Float32, 2, 3)),
     (ConvTranspose((3, 3), 3 => 2; stride=2), rand(Float32, 5, 5, 3, 1)),
     (StatefulRecurrentCell(RNNCell(3 => 5)), rand(Float32, 3, 2)),

test/helpers/compact_tests.jl

@@ -1,4 +1,4 @@
-@testitem "@compact" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "@compact" setup=[SharedTestSetup] tags=[:misc] begin
     using ComponentArrays, Zygote
 
     rng = StableRNG(12345)
@@ -439,7 +439,7 @@
     end
 end
 
-@testitem "@compact error checks" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "@compact error checks" setup=[SharedTestSetup] tags=[:misc] begin
     showerror(stdout, Lux.LuxCompactModelParsingException(""))
     println()
 

test/helpers/loss_tests.jl

@@ -1,4 +1,4 @@
-@testitem "LuxOps.xlogx & LuxOps.xlogy" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "LuxOps.xlogx & LuxOps.xlogy" setup=[SharedTestSetup] tags=[:misc] begin
     using ForwardDiff, Zygote, Enzyme
 
     @test iszero(LuxOps.xlogx(0))
@@ -55,7 +55,7 @@
     end
 end
 
-@testitem "Regression Loss" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Regression Loss" setup=[SharedTestSetup] tags=[:misc] begin
     using Zygote
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES
@@ -99,7 +99,7 @@ end
     end
 end
 
-@testitem "Classification Loss" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Classification Loss" setup=[SharedTestSetup] tags=[:misc] begin
     using OneHotArrays, Zygote
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES
@@ -283,7 +283,7 @@ end
     end
 end
 
-@testitem "Other Losses" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Other Losses" setup=[SharedTestSetup] tags=[:misc] begin
     using Zygote
 
     @testset "$mode" for (mode, aType, dev, ongpu) in MODES
@@ -404,7 +404,7 @@ end
     end
 end
 
-@testitem "Losses: Error Checks and Misc" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Losses: Error Checks and Misc" setup=[SharedTestSetup] tags=[:misc] begin
     @testset "Size Checks" begin
         @test_throws DimensionMismatch MSELoss()([1, 2], [1, 2, 3])
     end

test/helpers/size_propagator_test.jl

@@ -1,4 +1,4 @@
-@testitem "Size Propagator" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Size Propagator" setup=[SharedTestSetup] tags=[:misc] begin
     rng = StableRNG(12345)
 
     @testset "Simple Chain (LeNet)" begin

test/helpers/size_propagator_tests.jl

@@ -1,4 +1,4 @@
-@testitem "Size Propagator" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Size Propagator" setup=[SharedTestSetup] tags=[:misc] begin
     rng = StableRNG(12345)
 
     @testset "Simple Chain (LeNet)" begin

test/helpers/stateful_tests.jl

@@ -1,4 +1,4 @@
-@testitem "Simple Stateful Tests" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Simple Stateful Tests" setup=[SharedTestSetup] tags=[:misc] begin
     using Setfield, Zygote
 
     rng = StableRNG(12345)

test/helpers/training_tests.jl

@@ -1,4 +1,4 @@
-@testitem "TrainState" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "TrainState" setup=[SharedTestSetup] tags=[:misc] begin
     using Optimisers
 
     rng = StableRNG(12345)
@@ -19,7 +19,7 @@
     end
 end
 
-@testitem "AbstractADTypes" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "AbstractADTypes" setup=[SharedTestSetup] tags=[:misc] begin
     using ADTypes, Optimisers
 
     function _loss_function(model, ps, st, data)
@@ -50,7 +50,7 @@ end
     end
 end
 
-@testitem "Training API" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Training API" setup=[SharedTestSetup] tags=[:misc] begin
     using ADTypes, Optimisers
 
     mse = MSELoss()
@@ -125,7 +125,7 @@ end
     end
 end
 
-@testitem "Enzyme: Invalidate Cache on State Update" setup=[SharedTestSetup] tags=[:helpers] skip=:(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin
+@testitem "Enzyme: Invalidate Cache on State Update" setup=[SharedTestSetup] tags=[:misc] skip=:(using LuxTestUtils; !LuxTestUtils.ENZYME_TESTING_ENABLED) begin
     using ADTypes, Optimisers
 
     mse = MSELoss()
@@ -196,7 +196,7 @@ end
     @test hasfield(typeof(tstate_new2.cache.extras), :reverse)
 end
 
-@testitem "Compiled ReverseDiff" setup=[SharedTestSetup] tags=[:helpers] begin
+@testitem "Compiled ReverseDiff" setup=[SharedTestSetup] tags=[:misc] begin
     using ADTypes, Optimisers, ReverseDiff
 
     mse1 = MSELoss()

test/qa_tests.jl

@@ -1,4 +1,4 @@
-@testitem "Aqua: Quality Assurance" tags=[:others] begin
+@testitem "Aqua: Quality Assurance" tags=[:misc] begin
     using Aqua, ChainRulesCore, ForwardDiff
 
     Aqua.test_all(Lux; ambiguities=false, piracies=false)
@@ -10,7 +10,7 @@
     Aqua.test_piracies(Lux; treat_as_own=[Lux.outputsize])
 end
 
-@testitem "Explicit Imports: Quality Assurance" tags=[:others] begin
+@testitem "Explicit Imports: Quality Assurance" tags=[:misc] begin
     # Load all trigger packages
     import Lux, ComponentArrays, ReverseDiff, SimpleChains, Tracker, Zygote, Enzyme
     using ExplicitImports
@@ -30,7 +30,7 @@ end
 end
 
 # Some of the tests are flaky on prereleases
-@testitem "doctests: Quality Assurance" tags=[:others] begin
+@testitem "doctests: Quality Assurance" tags=[:misc] begin
     using Documenter
 
     doctestexpr = quote

test/runtests.jl

@@ -5,8 +5,8 @@ using InteractiveUtils, Hwloc
 
 const BACKEND_GROUP = lowercase(get(ENV, "BACKEND_GROUP", "all"))
 const ALL_LUX_TEST_GROUPS = [
-    "core_layers", "contrib", "helpers", "distributed", "normalize_layers",
-    "others", "autodiff", "recurrent_layers", "fluxcompat"]
+    "core_layers", "normalize_layers", "autodiff", "recurrent_layers", "misc"
+]
 
 Sys.iswindows() || push!(ALL_LUX_TEST_GROUPS, "reactant")
 
@@ -22,13 +22,12 @@ end
 const EXTRA_PKGS = Pkg.PackageSpec[]
 const EXTRA_DEV_PKGS = Pkg.PackageSpec[]
 
-if ("all" in LUX_TEST_GROUP || "distributed" in LUX_TEST_GROUP)
+if ("all" in LUX_TEST_GROUP || "misc" in LUX_TEST_GROUP)
     push!(EXTRA_PKGS, Pkg.PackageSpec("MPI"))
     (BACKEND_GROUP == "all" || BACKEND_GROUP == "cuda") &&
         push!(EXTRA_PKGS, Pkg.PackageSpec("NCCL"))
-end
-("all" in LUX_TEST_GROUP || "fluxcompat" in LUX_TEST_GROUP) &&
     push!(EXTRA_PKGS, Pkg.PackageSpec("Flux"))
+end
 
 if !Sys.iswindows()
     ("all" in LUX_TEST_GROUP || "reactant" in LUX_TEST_GROUP) &&
@@ -100,7 +99,7 @@ if "all" in LUX_TEST_GROUP || "core_layers" in LUX_TEST_GROUP
 end
 
 # Eltype Matching Tests
-if ("all" in LUX_TEST_GROUP || "eltype_match" in LUX_TEST_GROUP)
+if ("all" in LUX_TEST_GROUP || "misc" in LUX_TEST_GROUP)
     @testset "eltype_mismath_handling: $option" for option in (
         "none", "warn", "convert", "error")
         set_preferences!(Lux, "eltype_mismatch_handling" => option; force=true)
@@ -120,18 +119,23 @@ Lux.set_dispatch_doctor_preferences!(; luxcore="error", luxlib="error")
 const RETESTITEMS_NWORKERS = parse(
     Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 4))))
 
+const RETESTITEMS_NWORKER_THREADS = parse(
+    Int, get(ENV, "RETESTITEMS_NWORKER_THREADS",
+        string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1))))
+
 @testset "Lux.jl Tests" begin
     for (i, tag) in enumerate(LUX_TEST_GROUP)
-        (tag == "distributed" || tag == "eltype_match") && continue
+        tag == "misc" && continue
         @info "Running tests for group: [$(i)/$(length(LUX_TEST_GROUP))] $tag"
 
         ReTestItems.runtests(Lux; tags=(tag == "all" ? nothing : [Symbol(tag)]),
-            nworkers=RETESTITEMS_NWORKERS, testitem_timeout=2400)
+            nworkers=RETESTITEMS_NWORKERS, nworker_threads=RETESTITEMS_NWORKER_THREADS,
+            testitem_timeout=2400)
     end
 end
 
 # Distributed Tests
-if ("all" in LUX_TEST_GROUP || "distributed" in LUX_TEST_GROUP)
+if ("all" in LUX_TEST_GROUP || "misc" in LUX_TEST_GROUP)
     using MPI
 
     nprocs_str = get(ENV, "JULIA_MPI_TEST_NPROCS", "")