fix: minor test fixes (#98)

* fix: minor test fixes * chore: bump Zygote version
LuxDL · Jan 24, 2025 · 29c21ff · 29c21ff · avik-pal · Jan 24, 2025
1 parent f4e89a4
commit 29c21ff
Show file tree

Hide file tree

Showing 6 changed files with 12 additions and 6 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Boltz"
 uuid = "4544d5e4-abc5-4dea-817f-29e4c205d9c8"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.1.0"
+version = "1.1.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -70,5 +70,5 @@ Static = "1.1.1"
 Statistics = "1.10"
 Tracker = "0.2.34"
 WeightInitializers = "1"
-Zygote = "0.6.70"
+Zygote = "0.6.70, 0.7"
 julia = "1.10"
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -18,4 +18,4 @@ DynamicExpressions = "0.16, 0.17, 0.18, 0.19, 1"
 Literate = "2"
 Lux = "1"
 Random = "1.10"
-Zygote = "0.6.70"
+Zygote = "0.6.70, 0.7"
diff --git a/src/utils.jl b/src/utils.jl
@@ -4,7 +4,7 @@ using ForwardDiff: ForwardDiff
 using GPUArraysCore: AnyGPUArray
 using Statistics: mean
 
-using MLDataDevices: get_device_type, get_device, CPUDevice, CUDADevice
+using MLDataDevices: MLDataDevices, get_device_type, get_device, CPUDevice, CUDADevice
 
 is_extension_loaded(::Val) = false
 
@@ -88,4 +88,6 @@ struct DataTransferBarrier{V}
     val::V
 end
 
+MLDataDevices.isleaf(::DataTransferBarrier) = true
+
 end
diff --git a/test/Project.toml b/test/Project.toml
@@ -42,4 +42,4 @@ ReTestItems = "1.24.0"
 Reexport = "1.2.2"
 StableRNGs = "1.0.2"
 Test = "1.10"
-Zygote = "0.6.70"
+Zygote = "0.6.70, 0.7"
diff --git a/test/layer_tests.jl b/test/layer_tests.jl
@@ -23,7 +23,7 @@
 
                 __f = (x, ps) -> sum(abs2, first(model(x, ps, st)))
                 @test_gradients(__f, x, ps; atol=1e-3, rtol=1e-3,
-                    soft_fail=[AutoFiniteDiff()])
+                    soft_fail=[AutoFiniteDiff()], enzyme_set_runtime_activity=true)
             end
         end
     end

diff --git a/test/vision_tests.jl b/test/vision_tests.jl
@@ -105,6 +105,8 @@ end
 
     for (mode, aType, dev, ongpu) in MODES, depth in [18, 34, 50, 101, 152]
         @testset for pretrained in [false, true]
+            pretrained && pkgversion(Metalhead) > v"0.9.4" && continue
+
             model = Vision.ResNet(depth; pretrained)
             ps, st = Lux.setup(Random.default_rng(), model) |> dev
             st = Lux.testmode(st)
@@ -130,6 +132,7 @@ end
             (50, 32, 4), (101, 32, 8), (101, 64, 4), (152, 64, 4)]
             @testset for pretrained in [false, true]
                 depth == 152 && pretrained && continue
+                pretrained && pkgversion(Metalhead) > v"0.9.4" && continue
 
                 model = Vision.ResNeXt(depth; pretrained, cardinality, base_width)
                 ps, st = Lux.setup(Random.default_rng(), model) |> dev
@@ -155,6 +158,7 @@ end
     for (mode, aType, dev, ongpu) in MODES, depth in [50, 101, 152]
         @testset for pretrained in [false, true]
             depth == 152 && pretrained && continue
+            pretrained && pkgversion(Metalhead) > v"0.9.4" && continue
 
             model = Vision.WideResNet(depth; pretrained)
             ps, st = Lux.setup(Random.default_rng(), model) |> dev