test(LuxLib): migrate testing to v1.12 (#1633)

Author: avik-pal

.buildkite/testing_luxlib.yml

@@ -26,7 +26,7 @@ steps:
         matrix:
           setup:
             julia:
-              - "1.11"
+              - "1.12"
             group:
               - "common"
               - "normalization"

.github/workflows/CI_LuxLib.yml

@@ -40,7 +40,7 @@ jobs:
             loopvec: "false"
     uses: ./.github/workflows/CommonCI.yml
     with:
-      julia_version: "1.11"
+      julia_version: "1.12"
       project: "lib/LuxLib"
       test_args: "BACKEND_GROUP=cpu LUXLIB_TEST_GROUP=${{ matrix.test_group }} LUXLIB_BLAS_BACKEND=${{ matrix.blas_backend }} LUXLIB_LOAD_LOOPVEC=${{ matrix.loopvec }}"
 

lib/LuxLib/Project.toml

@@ -1,6 +1,6 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
-version = "1.15.1"
+version = "1.15.2"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]

lib/LuxLib/src/api/groupnorm.jl

@@ -40,19 +40,29 @@ end
 function assert_valid_groupnorm_arguments(
     x::AbstractArray{T,N}, scale, bias, groups
 ) where {T,N}
-    @assert length(scale) == length(bias) == size(x, N - 1) "Length of `scale` and `bias` must \
-                                                         be equal to the number of \
-                                                         channels ((N - 1) dim of the \
-                                                         input array)."
+    if length(scale) !== length(bias) || length(bias) !== size(x, N - 1)
+        throw(
+            ArgumentError(
+                "Length of `scale` and `bias` must be equal to the number of channels \
+                 ((N - 1) dim of the input array)."
+            ),
+        )
+    end
     assert_valid_groupnorm_arguments(x, nothing, nothing, groups)
     return nothing
 end
 
 function assert_valid_groupnorm_arguments(
     x::AbstractArray{T,N}, ::Nothing, ::Nothing, groups::Int
 ) where {T,N}
-    @assert size(x, N - 1) % groups == 0 "Number of channels $(size(x, N - 1)) must be \
-                                        divisible by the number of groups $groups."
+    if size(x, N - 1) % groups != 0
+        throw(
+            ArgumentError(
+                "Number of channels $(size(x, N - 1)) must be divisible by the number of \
+                 groups $groups."
+            ),
+        )
+    end
     return nothing
 end
 

lib/LuxLib/test/common_ops/activation_tests.jl

@@ -52,12 +52,14 @@
             @test_gradients(apply_act_fast, f, x; atol, rtol)
             @test_gradients(apply_act_fast2, f, x; atol, rtol)
 
-            ∂x1 = Zygote.gradient(apply_act, f, x)[2]
-            ∂x2 = Zygote.gradient(apply_act_fast, f, x)[2]
-            ∂x3 = Zygote.gradient(apply_act_fast2, f, x)[2]
+            if LuxTestUtils.ZYGOTE_TESTING_ENABLED[]
+                ∂x1 = Zygote.gradient(apply_act, f, x)[2]
+                ∂x2 = Zygote.gradient(apply_act_fast, f, x)[2]
+                ∂x3 = Zygote.gradient(apply_act_fast2, f, x)[2]
 
-            @test ∂x1 ≈ ∂x2 atol = atol rtol = rtol
-            @test ∂x1 ≈ ∂x3 atol = atol rtol = rtol
+                @test ∂x1 ≈ ∂x2 atol = atol rtol = rtol
+                @test ∂x1 ≈ ∂x3 atol = atol rtol = rtol
+            end
         end
     end
 end

lib/LuxLib/test/common_ops/attention_tests.jl

@@ -1,5 +1,5 @@
 @testitem "Scaled Dot Product Attention" tags = [:misc] setup = [SharedTestSetup] begin
-    using LuxLib, Reactant, NNlib, Random, MLDataDevices, Zygote, Enzyme, Statistics
+    using LuxLib, Reactant, NNlib, Random, MLDataDevices, Enzyme, Statistics
 
     @testset "$mode" for (mode, aType, ongpu, fp64) in MODES
         @testset "Different Batch Sizes" begin

lib/LuxLib/test/common_ops/bias_act_tests.jl

@@ -79,14 +79,16 @@
                 soft_fail=fp16 ? [AutoFiniteDiff()] : []
             )
 
-            ∂x1, ∂b1 = Zygote.gradient(__Fix1(bias_act_loss1, act), x, b)
-            ∂x2, ∂b2 = Zygote.gradient(__Fix1(bias_act_loss2, act), x, b)
-            ∂x3, ∂b3 = Zygote.gradient(__Fix1(bias_act_loss3, act), x, b)
-
-            @test ∂x1 ≈ ∂x2 atol = atol rtol = rtol
-            @test ∂x1 ≈ ∂x3 atol = atol rtol = rtol
-            @test ∂b1 ≈ ∂b2 atol = atol rtol = rtol
-            @test ∂b1 ≈ ∂b3 atol = atol rtol = rtol
+            if LuxTestUtils.ZYGOTE_TESTING_ENABLED[]
+                ∂x1, ∂b1 = Zygote.gradient(__Fix1(bias_act_loss1, act), x, b)
+                ∂x2, ∂b2 = Zygote.gradient(__Fix1(bias_act_loss2, act), x, b)
+                ∂x3, ∂b3 = Zygote.gradient(__Fix1(bias_act_loss3, act), x, b)
+
+                @test ∂x1 ≈ ∂x2 atol = atol rtol = rtol
+                @test ∂x1 ≈ ∂x3 atol = atol rtol = rtol
+                @test ∂b1 ≈ ∂b2 atol = atol rtol = rtol
+                @test ∂b1 ≈ ∂b3 atol = atol rtol = rtol
+            end
         end
     end
 end

lib/LuxLib/test/common_ops/conv_tests.jl

@@ -1,5 +1,5 @@
 @testsetup module ConvSetup
-using LuxLib, LuxTestUtils, Random, Test, Zygote, NNlib
+using LuxLib, LuxTestUtils, Random, Test, NNlib
 
 expand(_, i::Tuple) = i
 expand(N, i::Integer) = ntuple(_ -> i, N)

lib/LuxLib/test/common_ops/dense_tests.jl

@@ -39,16 +39,18 @@ function run_dense_testing(Tw, Tx, M, N, hasbias, activation, aType, mode, ongpu
     y_zyg = fused_dense_bias_activation(activation, w, x, bias)
     @test y_simple ≈ y_zyg atol = atol rtol = rtol
 
-    _, ∂w_true, ∂x_true, ∂b_true = Zygote.gradient(
-        sum ∘ dense_simple, activation, w, x, bias
-    )
-    _, ∂w_zyg, ∂x_zyg, ∂b_zyg = Zygote.gradient(
-        sum ∘ fused_dense_bias_activation, activation, w, x, bias
-    )
-    @test ∂w_true ≈ ∂w_zyg atol = atol rtol = rtol
-    @test ∂x_true ≈ ∂x_zyg atol = atol rtol = rtol
-    if bias !== nothing
-        @test ∂b_true ≈ ∂b_zyg atol = atol rtol = rtol
+    if LuxTestUtils.ZYGOTE_TESTING_ENABLED[]
+        _, ∂w_true, ∂x_true, ∂b_true = Zygote.gradient(
+            sum ∘ dense_simple, activation, w, x, bias
+        )
+        _, ∂w_zyg, ∂x_zyg, ∂b_zyg = Zygote.gradient(
+            sum ∘ fused_dense_bias_activation, activation, w, x, bias
+        )
+        @test ∂w_true ≈ ∂w_zyg atol = atol rtol = rtol
+        @test ∂x_true ≈ ∂x_zyg atol = atol rtol = rtol
+        if bias !== nothing
+            @test ∂b_true ≈ ∂b_zyg atol = atol rtol = rtol
+        end
     end
 end
 
@@ -201,35 +203,39 @@ end
                 b_enz,
             )
 
-            _, pb_f = Zygote.pullback(fused_dense_bias_activation, act, weight, x, b)
-            _, dweight_zyg, dx_zyg, db_zyg = pb_f(dy)
+            if LuxTestUtils.ZYGOTE_TESTING_ENABLED[]
+                _, pb_f = Zygote.pullback(fused_dense_bias_activation, act, weight, x, b)
+                _, dweight_zyg, dx_zyg, db_zyg = pb_f(dy)
 
-            @test dweight ≈ dweight_zyg atol = 1.0e-3 rtol = 1.0e-3
-            @test dx ≈ dx_zyg atol = 1.0e-3 rtol = 1.0e-3
-            if hasbias
-                @test db ≈ db_zyg atol = 1.0e-3 rtol = 1.0e-3
+                @test dweight ≈ dweight_zyg atol = 1.0e-3 rtol = 1.0e-3
+                @test dx ≈ dx_zyg atol = 1.0e-3 rtol = 1.0e-3
+                if hasbias
+                    @test db ≈ db_zyg atol = 1.0e-3 rtol = 1.0e-3
+                end
             end
 
             (act === identity && hasbias) || continue
 
-            dweight .= 0
-            dx .= 0
-            db .= 0
-            Enzyme.autodiff(
-                Reverse,
-                matmuladd!,
-                Duplicated(y, copy(dy)),
-                Duplicated(weight, dweight),
-                Duplicated(x, dx),
-                b_enz,
-            )
-
-            _, pb_f = Zygote.pullback(LuxLib.Impl.matmuladd, weight, x, b)
-            dweight_zyg, dx_zyg, db_zyg = pb_f(dy)
-
-            @test dweight ≈ dweight_zyg atol = 1.0e-3 rtol = 1.0e-3
-            @test dx ≈ dx_zyg atol = 1.0e-3 rtol = 1.0e-3
-            @test db ≈ db_zyg atol = 1.0e-3 rtol = 1.0e-3
+            if LuxTestUtils.ZYGOTE_TESTING_ENABLED[]
+                dweight .= 0
+                dx .= 0
+                db .= 0
+                Enzyme.autodiff(
+                    Reverse,
+                    matmuladd!,
+                    Duplicated(y, copy(dy)),
+                    Duplicated(weight, dweight),
+                    Duplicated(x, dx),
+                    b_enz,
+                )
+
+                _, pb_f = Zygote.pullback(LuxLib.Impl.matmuladd, weight, x, b)
+                dweight_zyg, dx_zyg, db_zyg = pb_f(dy)
+
+                @test dweight ≈ dweight_zyg atol = 1.0e-3 rtol = 1.0e-3
+                @test dx ≈ dx_zyg atol = 1.0e-3 rtol = 1.0e-3
+                @test db ≈ db_zyg atol = 1.0e-3 rtol = 1.0e-3
+            end
         end
     end
 end

lib/LuxLib/test/normalization/batchnorm_tests.jl

@@ -73,7 +73,7 @@ function run_batchnorm_testing(gen_f, T, sz, training, affine, track_stats, act,
     end
 
     # Check the rrules
-    if is_training(training)
+    if is_training(training) && LuxTestUtils.ZYGOTE_TESTING_ENABLED[]
         _f =
             (args...) ->
                 sum(first(batchnorm(args..., rm, rv, training, act, T(0.9), epsilon)))