feat: allow setting target_modules at runtime

LuxDL · Jul 12, 2024 · d24bae7 · d24bae7 · avik-pal · Jul 12, 2024
2 parents bae193e + ac4a567
commit d24bae7
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,6 +2,7 @@
 *.jl.*.cov
 *.jl.mem
 /Manifest.toml
+Manifest-v*.toml
 /deps/deps.jl
 /docs/build
 /docs/Manifest.toml

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxTestUtils"
 uuid = "ac9de150-d08f-4546-94fb-7472b5760531"
 authors = ["Avik Pal <[email protected]>"]
-version = "0.1.16"
+version = "0.1.17"
 
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"

diff --git a/src/LuxTestUtils.jl b/src/LuxTestUtils.jl
@@ -3,7 +3,17 @@ module LuxTestUtils
 using ComponentArrays, Optimisers, Preferences, LuxCore, LuxDeviceUtils, Test
 using ForwardDiff, ReverseDiff, Tracker, Zygote, FiniteDifferences
 
-const JET_TARGET_MODULES = @load_preference("target_modules", nothing)
+const JET_TARGET_MODULES = Ref{Union{Nothing, Vector{String}}}(nothing)
+
+function __init__()
+    JET_TARGET_MODULES[] = @load_preference("target_modules", nothing)
+end
+
+function jet_target_modules!(list::Vector{String})
+    JET_TARGET_MODULES[] = list
+    @info "JET_TARGET_MODULES set to $list"
+    return list
+end
 
 # JET Testing
 try