fix: bad rebase

LuxDL · Aug 18, 2024 · 33a2d75 · 33a2d75
1 parent 803a43e
commit 33a2d75
Show file tree

Hide file tree

Showing 11 changed files with 8 additions and 302 deletions.
diff --git a/ext/LuxMPIExt.jl b/ext/LuxMPIExt.jl
@@ -1,9 +1,5 @@
 module LuxMPIExt
 
-using Lux: MPIBackend, NCCLBackend, DistributedUtils, __unwrap_val, MPI_CUDA_AWARE,
-           MPI_ROCM_AWARE
-using MLDataDevices: AbstractDevice, CUDADevice, AMDGPUDevice, cpu_device, set_device!,
-                     functional
 using MPI: MPI
 
 using Lux: Lux, MPIBackend, NCCLBackend, DistributedUtils, Utils

diff --git a/ext/LuxReverseDiffExt/LuxReverseDiffExt.jl b/ext/LuxReverseDiffExt/LuxReverseDiffExt.jl
@@ -10,11 +10,6 @@ using MLDataDevices: CPUDevice
 using ReverseDiff: ReverseDiff, ForwardExecutor, ReverseExecutor, TrackedArray, TrackedReal,
                    @grad_from_chainrules
 
-using Lux: Lux, Utils
-using Lux.Training: TrainingBackendCache, TrainState
-using LuxCore: LuxCore
-using MLDataDevices: CPUDevice
-
 include("utils.jl")
 include("rules.jl")
 include("training.jl")

diff --git a/ext/LuxTrackerExt/LuxTrackerExt.jl b/ext/LuxTrackerExt/LuxTrackerExt.jl
diff --git a/ext/LuxTrackerExt/rules.jl b/ext/LuxTrackerExt/rules.jl
diff --git a/ext/LuxTrackerExt/training.jl b/ext/LuxTrackerExt/training.jl
diff --git a/ext/LuxTrackerExt/utils.jl b/ext/LuxTrackerExt/utils.jl
diff --git a/ext/LuxZygoteExt/LuxZygoteExt.jl b/ext/LuxZygoteExt/LuxZygoteExt.jl
@@ -4,8 +4,6 @@ using ArgCheck: @argcheck
 using ADTypes: AutoZygote
 using ChainRulesCore: ChainRulesCore
 using ForwardDiff: ForwardDiff
-using Lux: Lux
-using MLDataDevices: get_device_type, CPUDevice
 using Setfield: @set!
 using Zygote: Zygote
 

diff --git a/src/contrib/map.jl b/src/contrib/map.jl
@@ -39,7 +39,7 @@ true
 ```
 """
 macro layer_map(f, l, ps, st)
-    quote
+    return quote
         layer_map($(esc(f)), $(esc(l)), $(esc(ps)), $(esc(st)), $(Meta.quot(l)))
     end
 end

diff --git a/src/helpers/nested_ad.jl b/src/helpers/nested_ad.jl
diff --git a/test/contrib/share_parameters_tests.jl b/test/contrib/share_parameters_tests.jl
@@ -17,9 +17,9 @@
         @test ps_1.d3.bias == ps_1.d2.l1.bias
 
         ps_new_1 = (; weight=randn(rng, Float32, 4, 2), bias=randn(rng, Float32, 4)) |>
-                   device
+                   dev
         ps_new_2 = (; weight=randn(rng, Float32, 2, 4), bias=randn(rng, Float32, 2)) |>
-                   device
+                   dev
 
         ps_2 = Lux.Experimental.share_parameters(ps, sharing, (ps_new_1, ps_new_2))
 
@@ -29,7 +29,7 @@
         @test ps_2.d3.bias == ps_new_2.bias == ps_2.d2.l1.bias
 
         # Mix in ComponentArray
-        ps_new_ca_1 = ComponentArray(ps_new_1 |> CPUDevice()) |> device
+        ps_new_ca_1 = ComponentArray(ps_new_1 |> CPUDevice()) |> dev
 
         ps_3 = Lux.Experimental.share_parameters(ps, sharing, (ps_new_ca_1, ps_new_2))
 
@@ -47,14 +47,14 @@
 
         # Parameter Structure Mismatch
         ps_new_1 = (; weight=randn(rng, Float32, 2, 4), bias=randn(rng, Float32, 4)) |>
-                   device
+                   dev
         ps_new_2 = (; weight=randn(rng, Float32, 2, 4), bias=randn(rng, Float32, 2)) |>
-                   device
+                   dev
 
         @test_throws ArgumentError Lux.Experimental.share_parameters(
             ps, sharing, (ps_new_1, ps_new_2))
 
-        ps_new_ca_1 = ComponentArray(ps_new_1 |> CPUDevice()) |> device
+        ps_new_ca_1 = ComponentArray(ps_new_1 |> CPUDevice()) |> dev
 
         @test_throws ArgumentError Lux.Experimental.share_parameters(
             ps, sharing, (ps_new_ca_1, ps_new_2))

diff --git a/test/utils_tests.jl b/test/utils_tests.jl
@@ -127,7 +127,7 @@ end
 @testitem "FP Conversions" setup=[SharedTestSetup] tags=[:others] begin
     rng = StableRNG(12345)
 
-    @testset "$mode" for (mode, aType, device, ongpu) in MODES
+    @testset "$mode" for (mode, aType, dev, ongpu) in MODES
         model = Chain(
             Dense(1 => 16, relu), Chain(Dense(16 => 1), Dense(1 => 1)), BatchNorm(1))
-Original file line number
+Diff line change
@@ Expand Up / @@ -39,7 +39,7 @@ true @@
     ```
     """
     macro layer_map(f, l, ps, st)
-        quote
+        return quote
             layer_map($(esc(f)), $(esc(l)), $(esc(ps)), $(esc(st)), $(Meta.quot(l)))
         end
     end
@@ Expand Down @@