LuxDL · avik-pal · Feb 3, 2026 · Feb 3, 2026 · Feb 3, 2026 · Feb 3, 2026
diff --git a/.buildkite/testing.yml b/.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:

diff --git a/.github/workflows/CI.yml b/.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"
diff --git a/Project.toml b/Project.toml
@@ -113,7 +113,7 @@ Enzyme = "0.13.120"
 EnzymeCore = "0.8.16"
 FastClosures = "0.3.2"
 Flux = "0.16.3"
-ForwardDiff = "0.10.36, 1"
+ForwardDiff = "1"
 FunctionWrappers = "1.1.3"
 Functors = "0.5"
 GPUArrays = "11"
@@ -129,7 +129,7 @@ MacroTools = "0.5.13"
 Markdown = "1.10"
 Mooncake = "0.5"
 NCCL = "0.1.2"
-NNlib = "0.9.27"
+NNlib = "0.9.33"
 Optimisers = "0.4.6"
 PrecompileTools = "1.2.1"
 Preferences = "1.4.3"

diff --git a/README.md b/README.md
@@ -271,51 +271,44 @@ format(".")
 
 The full test of `Lux.jl` takes a long time, here's how to test a portion of the code.
 
-For each `@testitem`, there are corresponding `tags`, for example:
+Tests are organized by directories, where each directory contains test files with `@testset`
+blocks. For example, tests for `SkipConnection` are in `test/core_layers/containers_tests.jl`.
 
-```julia
-@testitem "SkipConnection" setup=[SharedTestSetup] tags=[:core_layers]
-```
+#### Running a Specific Test File
 
-For example, let's consider the tests for `SkipConnection`:
+The easiest way to run a specific test is to directly activate the test directory and
+include the test file:
 
 ```julia
-@testitem "SkipConnection" setup=[SharedTestSetup] tags=[:core_layers] begin
-    ...
-end
-```
-
-We can test the group to which `SkipConnection` belongs by testing `core_layers`.
-To do so set the `LUX_TEST_GROUP` environment variable, or rename the tag to
-further narrow the test scope:
+# From the Lux.jl root directory
+using Pkg
+Pkg.activate("test")
 
-```shell
-export LUX_TEST_GROUP="core_layers"
+# Run a specific test file
+include("test/core_layers/containers_tests.jl")
 ```
 
-Or directly modify the default test tag in `runtests.jl`:
+This approach allows you to quickly iterate on specific tests without running the entire
+test suite.
 
-```julia
-# const LUX_TEST_GROUP = lowercase(get(ENV, "LUX_TEST_GROUP", "all"))
-const LUX_TEST_GROUP = lowercase(get(ENV, "LUX_TEST_GROUP", "core_layers"))
-```
+See [ParallelTestRunners.jl](https://github.com/JuliaTesting/ParallelTestRunner.jl) for
+details on executing specific groups of tests.
 
-But be sure to restore the default value "all" before submitting the code.
+#### Running Test Groups via CI
 
-Furthermore if you want to run a specific test based on the name of the testset, you can
-use [TestEnv.jl](https://github.com/JuliaTesting/TestEnv.jl) as follows. Start with activating the Lux environment and then run the following:
+To run a specific group of tests via the test runner, you can pass the directory name as a
+positional argument:
 
-```julia
-using TestEnv; TestEnv.activate(); using ReTestItems;
-
-# Assuming you are in the main directory of Lux
-ReTestItems.runtests("tests/"; name = "NAME OF THE TEST")
+```shell
+julia --project -e 'using Pkg; Pkg.test(test_args=["core_layers"])'
 ```
 
-For the `SkipConnection` tests that would be:
+#### Running All Tests
 
-```julia
-ReTestItems.runtests("tests/"; name = "SkipConnection")
+To run the full test suite:
+
+```shell
+julia --project -e 'using Pkg; Pkg.test()'
 ```
 
 ### 📖 Documentation

diff --git a/lib/LuxLib/Project.toml b/lib/LuxLib/Project.toml
@@ -1,6 +1,6 @@
 name = "LuxLib"
 uuid = "82251201-b29d-42c6-8e01-566dec8acb11"
-version = "1.15.3"
+version = "1.15.4"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]
@@ -68,7 +68,7 @@ TrackerExt = "Tracker"
 cuDNNExt = ["CUDA", "cuDNN"]
 
 [compat]
-AMDGPU = "0.9.6, 1, 2"
+AMDGPU = "2.2"
 AppleAccelerate = "0.4, 0.5"
 ArrayInterface = "7.17.1"
 BLISBLAS = "0.1, 0.2"
@@ -88,7 +88,7 @@ LuxCore = "1.5"
 MKL = "0.7, 0.8, 0.9"
 MLDataDevices = "1.17.1"
 Markdown = "1.10"
-NNlib = "0.9.27"
+NNlib = "0.9.33"
 Octavian = "0.3.28"
 OneHotArrays = "0.2.5"
 Polyester = "0.7.18"

diff --git a/lib/LuxLib/test/Project.toml b/lib/LuxLib/test/Project.toml
@@ -60,7 +60,7 @@ LuxTestUtils = "2"
 MKL = "0.7, 0.8, 0.9"
 MLDataDevices = "1.17"
 Mooncake = "0.5"
-NNlib = "0.9.27"
+NNlib = "0.9.33"
 Octavian = "0.3.28"
 ParallelTestRunner = "2.1"
 Pkg = "1.10"

diff --git a/lib/MLDataDevices/Project.toml b/lib/MLDataDevices/Project.toml
@@ -1,7 +1,7 @@
 name = "MLDataDevices"
 uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.17.4"
+version = "1.17.5"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -53,7 +53,7 @@ oneAPIExt = ["GPUArrays", "oneAPI"]
 OpenCLExt = ["GPUArrays", "OpenCL"]
 
 [compat]
-AMDGPU = "1, 2"
+AMDGPU = "2.2"
 Adapt = "4.4"
 CUDA = "5.8"
 ChainRules = "1.51"

diff --git a/lib/WeightInitializers/Project.toml b/lib/WeightInitializers/Project.toml
@@ -1,7 +1,7 @@
 name = "WeightInitializers"
 uuid = "d49dbf32-c5c2-4618-8acc-27bb2598ef2d"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.3.1"
+version = "1.3.2"
 
 [deps]
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
@@ -26,7 +26,7 @@ GPUArraysExt = "GPUArrays"
 ReactantExt = "Reactant"
 
 [compat]
-AMDGPU = "0.9.6, 1, 2"
+AMDGPU = "2.2"
 CUDA = "5.8"
 ChainRulesCore = "1.25.1"
 ConcreteStructs = "0.2.3"

diff --git a/test/Project.toml b/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,13 +27,12 @@ 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"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -41,6 +42,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 = ".."}
@@ -50,6 +52,7 @@ LuxTestUtils = {path = "../lib/LuxTestUtils"}
 MLDataDevices = {path = "../lib/MLDataDevices"}
 
 [compat]
+AMDGPU = "2.2"
 ADTypes = "1.19"
 Adapt = "4"
 Aqua = "0.8.4"
@@ -60,9 +63,8 @@ DispatchDoctor = "0.4.12"
 Documenter = "1.4"
 Enzyme = "0.13.120"
 ExplicitImports = "1.14.2"
-ForwardDiff = "0.10.36, 1"
+ForwardDiff = "1"
 Functors = "0.5"
-InteractiveUtils = "<0.0.1, 1"
 LinearAlgebra = "1.10"
 Logging = "1.10"
 LoopVectorization = "0.12.171"
@@ -71,17 +73,16 @@ LuxLib = "1.15"
 LuxTestUtils = "2"
 MLDataDevices = "1.17"
 Mooncake = "0.4.138, 0.5"
-NNlib = "0.9.27"
+NNlib = "0.9.33"
 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"
 Setfield = "1.1.1"
 StableRNGs = "1.0.2"