test(Lux): migrate to ParallelTestRunner

Author: avik-pal

.buildkite/testing.yml

@@ -6,7 +6,7 @@ steps:
           - JuliaCI/julia#v1:
               version: "{{matrix.julia}}"
           - JuliaCI/julia-test#v1:
-              test_args: "BACKEND_GROUP=CUDA LUX_TEST_GROUP={{matrix.group}}"
+              test_args: "--BACKEND_GROUP=CUDA {{matrix.group}}"
           - JuliaCI/julia-coverage#v1:
               codecov: true
               dirs:
@@ -31,13 +31,14 @@ steps:
             julia:
               - "1.12"
             group:
+              - "autodiff"
+              - "contrib"
+              - "helpers"
               - "core_layers"
-              - "normalize_layers"
               - "recurrent_layers"
-              - "autodiff"
               - "misc"
               - "reactant"
-              - "extras"
+              - "others"
 
   # - group: ":julia: (Lux) AMD GPU"
   #   steps:
@@ -46,7 +47,7 @@ steps:
   #         - JuliaCI/julia#v1:
   #             version: "{{matrix.julia}}"
   #         - JuliaCI/julia-test#v1:
-  #             test_args: "BACKEND_GROUP=AMDGPU"
+  #             test_args: "--BACKEND_GROUP=AMDGPU"
   #         - JuliaCI/julia-coverage#v1:
   #             codecov: true
   #             dirs:

.github/workflows/CI.yml

@@ -39,18 +39,17 @@ jobs:
         os:
           - ubuntu-latest
         test_group:
+          - "autodiff"
           - "core_layers"
-          - "normalize_layers"
           - "recurrent_layers"
-          - "autodiff"
           - "misc"
           - "reactant"
-          - "extras"
+          - "others"
     uses: ./.github/workflows/CommonCI.yml
     with:
       julia_version: ${{ matrix.version }}
       project: "."
-      test_args: "BACKEND_GROUP=CPU LUX_TEST_GROUP=${{ matrix.test_group }}"
+      test_args: "--BACKEND_GROUP=CPU ${{ matrix.test_group }}"
       os: ${{ matrix.os }}
       local_dependencies: "lib/MLDataDevices,lib/WeightInitializers,lib/LuxLib,lib/LuxCore"
       local_test_dependencies: "lib/MLDataDevices,lib/LuxTestUtils,lib/LuxLib,lib/LuxCore"
@@ -60,18 +59,17 @@ jobs:
       fail-fast: false
       matrix:
         test_group:
+          - "autodiff"
           - "core_layers"
-          - "normalize_layers"
           - "recurrent_layers"
-          - "autodiff"
           - "misc"
           - "reactant"
-          - "extras"
+          - "others"
     uses: ./.github/workflows/CommonCI.yml
     with:
       julia_version: "1.11"
       project: "."
       downgrade_testing: true
-      test_args: "BACKEND_GROUP=CPU LUX_TEST_GROUP=${{ matrix.test_group }}"
+      test_args: "--BACKEND_GROUP=CPU ${{ matrix.test_group }}"
       local_dependencies: "lib/MLDataDevices,lib/WeightInitializers,lib/LuxLib,lib/LuxCore"
       local_test_dependencies: "lib/MLDataDevices,lib/LuxTestUtils,lib/LuxLib,lib/LuxCore"

test/Project.toml

@@ -1,8 +1,10 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CPUSummary = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
@@ -11,11 +13,11 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
@@ -25,11 +27,11 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Octavian = "6fd5a793-0b7e-452c-907f-f8bfe9c57db4"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -41,6 +43,7 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd"
 
 [sources]
 Lux = {path = ".."}
@@ -62,7 +65,6 @@ Enzyme = "0.13.120"
 ExplicitImports = "1.14.2"
 ForwardDiff = "0.10.36, 1"
 Functors = "0.5"
-InteractiveUtils = "<0.0.1, 1"
 LinearAlgebra = "1.10"
 Logging = "1.10"
 LoopVectorization = "0.12.171"
@@ -75,11 +77,11 @@ NNlib = "0.9.27"
 Octavian = "0.3.28"
 OneHotArrays = "0.2.5"
 Optimisers = "0.4.6"
+ParallelTestRunner = "2.1"
 Pkg = "1.10"
 Preferences = "1.4.3"
 PythonCall = "0.9"
 Random = "1.10"
-ReTestItems = "1.24.0"
 Reactant = "0.2.205"
 Reexport = "1.2.2"
 ReverseDiff = "1.15.3"

test/autodiff/batched_jacobian_tests.jl

@@ -0,0 +1,99 @@
+using ComponentArrays, ForwardDiff, Zygote, Tracker, ReverseDiff
+
+include("../shared_testsetup.jl")
+
+@testset "Batched Jacobian" begin
+    rng = StableRNG(12345)
+
+    @testset "$mode" for (mode, aType, dev, ongpu) in MODES
+        models = (
+            Chain(
+                Conv((3, 3), 2 => 4, gelu; pad=SamePad()),
+                Conv((3, 3), 4 => 2, gelu; pad=SamePad()),
+                FlattenLayer(),
+                Dense(18 => 2),
+            ),
+            Chain(Dense(2, 4, gelu), Dense(4, 2)),
+        )
+        Xs = (aType(randn(rng, Float32, 3, 3, 2, 4)), aType(randn(rng, Float32, 2, 4)))
+
+        for (model, X) in zip(models, Xs)
+            ps, st = dev(Lux.setup(rng, model))
+            smodel = StatefulLuxLayer(model, ps, st)
+
+            J1 = allow_unstable() do
+                ForwardDiff.jacobian(smodel, X)
+            end
+
+            @testset "$(backend)" for backend in (AutoZygote(), AutoForwardDiff())
+                J2 = allow_unstable() do
+                    batched_jacobian(smodel, backend, X)
+                end
+                J2_mat =
+                    mapreduce(
+                        Base.Fix1(Lux.AutoDiffInternalImpl.batched_row, J2),
+                        hcat,
+                        1:(size(J2, 1) * size(J2, 3)),
+                    )'
+
+                @test Array(J1) ≈ Array(J2_mat) atol = 1.0e-3 rtol = 1.0e-3
+
+                ps = dev(ComponentArray(cpu_device()(ps)))
+
+                smodel = StatefulLuxLayer(model, ps, st)
+
+                J3 = allow_unstable() do
+                    batched_jacobian(smodel, backend, X)
+                end
+
+                @test Array(J2) ≈ Array(J3) atol = 1.0e-3 rtol = 1.0e-3
+            end
+        end
+
+        @testset "Issue #636 Chunksize Specialization" begin
+            for N in (2, 4, 8, 11, 12, 50, 51), backend in (AutoZygote(), AutoForwardDiff())
+                model = @compact(; potential=Dense(N => N, gelu), backend=backend) do x
+                    @return allow_unstable() do
+                        batched_jacobian(potential, backend, x)
+                    end
+                end
+
+                ps, st = dev(Lux.setup(Random.default_rng(), model))
+
+                x = dev(randn(Float32, N, 3))
+                @test first(model(x, ps, st)) isa AbstractArray{<:Any,3}
+            end
+        end
+
+        @testset "Simple Batched Jacobian" begin
+            # Without any Lux bs just plain old batched jacobian
+            ftest(x) = x .^ 2
+            x = dev(reshape(Float32.(1:6), 2, 3))
+
+            Jx_true = zeros(Float32, 2, 2, 3)
+            Jx_true[1, 1, 1] = 2
+            Jx_true[2, 2, 1] = 4
+            Jx_true[1, 1, 2] = 6
+            Jx_true[2, 2, 2] = 8
+            Jx_true[1, 1, 3] = 10
+            Jx_true[2, 2, 3] = 12
+            Jx_true = dev(Jx_true)
+
+            Jx_fdiff = allow_unstable() do
+                batched_jacobian(ftest, AutoForwardDiff(), x)
+            end
+            @test Jx_fdiff ≈ Jx_true
+
+            Jx_zygote = allow_unstable() do
+                batched_jacobian(ftest, AutoZygote(), x)
+            end
+            @test Jx_zygote ≈ Jx_true
+
+            fincorrect(x) = x[:, 1]
+            x = dev(reshape(Float32.(1:6), 2, 3))
+
+            @test_throws AssertionError batched_jacobian(fincorrect, AutoForwardDiff(), x)
+            @test_throws AssertionError batched_jacobian(fincorrect, AutoZygote(), x)
+        end
+    end
+end

test/enzyme_tests.jl test/autodiff/enzyme_tests.jltest/enzyme_tests.jl renamed to test/autodiff/enzyme_tests.jl

@@ -1,9 +1,12 @@
 # With LuxTestUtils 1.1 we have inbuilt Enzyme.jl support in `test_gradients`. We should be
 # able to remove this, but this file is still helpful to catch errors in a localized way.
-@testsetup module EnzymeTestSetup
+
 using LuxTestUtils, Enzyme, Zygote, Test
 using Lux, NNlib
 using LuxTestUtils: check_approx
+using ComponentArrays
+
+include("../shared_testsetup.jl")
 
 generic_loss_function(model, x, ps, st) = sum(abs2, first(model(x, ps, st)))
 
@@ -108,49 +111,27 @@ if VERSION < v"1.11-"
     )
 end
 
-export generic_loss_function,
-    compute_enzyme_gradient, compute_zygote_gradient, test_enzyme_gradients, MODELS_LIST
-end
-
-@testitem "Enzyme Integration" setup = [EnzymeTestSetup, SharedTestSetup] tags = [
-    :autodiff, :enzyme
-] timeout = 3600 skip = :(using LuxTestUtils;
-!LuxTestUtils.ENZYME_TESTING_ENABLED[]) begin
-    rng = StableRNG(12345)
-
-    @testset "$mode" for (mode, aType, dev, ongpu) in MODES
-        ongpu && continue
-
-        @testset "[$(i)] $(nameof(typeof(model)))" for (i, (model, x)) in
-                                                       enumerate(MODELS_LIST)
-            display(model)
-
-            ps, st = dev(Lux.setup(rng, model))
-            x = aType(x)
-            test_enzyme_gradients(model, x, ps, st)
-        end
-    end
-end
-
-@testitem "Enzyme Integration ComponentArray" setup = [EnzymeTestSetup, SharedTestSetup] timeout =
-    3600 tags = [:autodiff, :enzyme] skip = :(using LuxTestUtils;
-!LuxTestUtils.ENZYME_TESTING_ENABLED[]) begin
-    using ComponentArrays
+@testset "Enzyme Integration" begin
+    if LuxTestUtils.ENZYME_TESTING_ENABLED[]
+        rng = StableRNG(12345)
 
-    rng = StableRNG(12345)
+        @testset "$mode" for (mode, aType, dev, ongpu) in MODES
+            ongpu && continue
 
-    @testset "$mode" for (mode, aType, dev, ongpu) in MODES
-        ongpu && continue
+            @testset "[$(i)] $(nameof(typeof(model)))" for (i, (model, x)) in
+                                                           enumerate(MODELS_LIST)
+                display(model)
 
-        @testset "[$(i)] $(nameof(typeof(model)))" for (i, (model, x)) in
-                                                       enumerate(MODELS_LIST)
-            display(model)
+                ps, st = dev(Lux.setup(rng, model))
+                x = aType(x)
+                test_enzyme_gradients(model, x, ps, st)
 
-            ps, st = Lux.setup(rng, model)
-            ps = ComponentArray(ps)
-            st = dev(st)
-            x = aType(x)
-            test_enzyme_gradients(model, x, ps, st)
+                @testset "ComponentArray" begin
+                    ps_ca = ComponentArray(ps)
+                    st_ca = dev(st)
+                    test_enzyme_gradients(model, x, ps_ca, st_ca)
+                end
+            end
         end
     end
 end