From 425deaf468eafb4bd181dd855c29f8d7b4ea6d1b Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 20 Sep 2025 22:17:39 -0400
Subject: [PATCH 01/16] ci: use a shared workflow

---
 .github/workflows/CI.yml | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index b60f801cbc..0acd10ca52 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -67,24 +67,24 @@ jobs:
       runtime: ${{ matrix.runtime }}
       test_group: ${{ matrix.test_group }}
 
-  # This has been broken for a while, originating from CUDA.jl
-  # test-assertions:
-  #   strategy:
-  #     fail-fast: false
-  #     matrix:
-  #       version:
-  #         - "1.10"
-  #       test_group:
-  #         - core
-  #         - neural_networks
-  #         - integration
-  #   uses: ./.github/workflows/CommonCI.yml
-  #   with:
-  #     julia_version: ${{ matrix.version }}
-  #     os: "ubuntu-24.04"
-  #     runtime: "both"
-  #     assertions: true
-  #     test_group: ${{ matrix.test_group }}
+  test-assertions:
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - "1.10"
+          - "1.11"
+        test_group:
+          - core
+          - neural_networks
+          - integration
+    uses: ./.github/workflows/CommonCI.yml
+    with:
+      julia_version: ${{ matrix.version }}
+      os: "ubuntu-24.04"
+      runtime: "both"
+      assertions: true
+      test_group: ${{ matrix.test_group }}
 
   downgrade:
     strategy:

From 9546706f8c4d8f7b530b30985798216d3c12c410 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 11:05:59 -0400
Subject: [PATCH 02/16] ci: disable assertions for now

---
 .github/workflows/CI.yml | 36 ++++++++++++++++++------------------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 0acd10ca52..9d38b2e705 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -67,24 +67,24 @@ jobs:
       runtime: ${{ matrix.runtime }}
       test_group: ${{ matrix.test_group }}
 
-  test-assertions:
-    strategy:
-      fail-fast: false
-      matrix:
-        version:
-          - "1.10"
-          - "1.11"
-        test_group:
-          - core
-          - neural_networks
-          - integration
-    uses: ./.github/workflows/CommonCI.yml
-    with:
-      julia_version: ${{ matrix.version }}
-      os: "ubuntu-24.04"
-      runtime: "both"
-      assertions: true
-      test_group: ${{ matrix.test_group }}
+  # This has been broken for a while, originating from Enzyme.jl
+  # test-assertions:
+  #   strategy:
+  #     fail-fast: false
+  #     matrix:
+  #       version:
+  #         - "1.10"
+  #       test_group:
+  #         - core
+  #         - neural_networks
+  #         - integration
+  #   uses: ./.github/workflows/CommonCI.yml
+  #   with:
+  #     julia_version: ${{ matrix.version }}
+  #     os: "ubuntu-24.04"
+  #     runtime: "both"
+  #     assertions: true
+  #     test_group: ${{ matrix.test_group }}
 
   downgrade:
     strategy:

From d2d61f319a9d9dc67a108694717f0c4a2cf80f7c Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 14:07:46 -0400
Subject: [PATCH 03/16] ci: fix cache name

---
 .github/workflows/CI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 9d38b2e705..b60f801cbc 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -67,7 +67,7 @@ jobs:
       runtime: ${{ matrix.runtime }}
       test_group: ${{ matrix.test_group }}
 
-  # This has been broken for a while, originating from Enzyme.jl
+  # This has been broken for a while, originating from CUDA.jl
   # test-assertions:
   #   strategy:
   #     fail-fast: false

From 2e2b928c37ceadbe6012686b170a4865b5db30cf Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sat, 20 Sep 2025 19:42:53 -0500
Subject: [PATCH 04/16] refactor: move tests into separate projects

---
 test/Project.toml                             |  24 -----
 .../special_functions.jl                      |   0
 test/{ => core}/autodiff.jl                   |   0
 test/{ => core}/basic.jl                      |   0
 test/{ => core}/batching.jl                   |   0
 test/{ => core}/bcast.jl                      |   0
 test/{ => core}/buffer_donation.jl            |   0
 test/{ => core}/closure.jl                    |   0
 test/{ => core}/compile.jl                    |   0
 test/{ => core}/complex.jl                    |   0
 test/{ => core}/config.jl                     |   0
 test/{ => core}/constructor.jl                |   0
 test/{ => core}/control_flow.jl               |   0
 test/{ => core}/custom_number_types.jl        |   0
 test/{ => core}/indexing.jl                   |   0
 test/{ => core}/ir.jl                         |   0
 test/{ => core}/layout.jl                     |   0
 test/{ => core}/ops.jl                        |   0
 test/{ => core}/qa.jl                         |   0
 test/{ => core}/sorting.jl                    |   0
 test/{ => core}/struct.jl                     |   0
 test/{ => core}/tracing.jl                    |   0
 test/{ => core}/wrapped_arrays.jl             |   0
 test/{ => distributed}/cluster_detector.jl    |   0
 test/{ => distributed}/optimize_comm.jl       |   0
 test/{ => distributed}/sharding.jl            |   0
 test/integration/CUDA/Project.toml            |  11 ++
 test/integration/{ => CUDA}/cuda.jl           |   0
 .../{ => CUDA}/kernelabstractions.jl          |   0
 test/integration/CUDA/runtests.jl             |  11 ++
 test/integration/FFT/Project.toml             |  10 ++
 test/integration/{fft.jl => FFT/runtests.jl}  |   0
 test/integration/FancyArrayTypes/Project.toml |  13 +++
 .../{ => FancyArrayTypes}/fillarrays.jl       |   0
 .../{ => FancyArrayTypes}/offsetarrays.jl     |   5 +-
 .../{ => FancyArrayTypes}/onehotarrays.jl     |   0
 test/integration/FancyArrayTypes/runtests.jl  |  15 +++
 test/integration/Flux/Project.toml            |  10 ++
 test/integration/Flux/runtests.jl             |  24 +++++
 test/integration/Lux/Project.toml             |  17 +++
 test/{nn => integration/Lux}/lux.jl           |   0
 test/{nn => integration/Lux}/luxlib.jl        |   0
 test/integration/Lux/runtests.jl              |  11 ++
 test/integration/NNlib/Project.toml           |  11 ++
 .../NNlib/runtests.jl}                        |   3 +-
 test/integration/Optimisers/Project.toml      |   4 +
 .../{optimisers.jl => Optimisers/runtests.jl} |   0
 test/integration/Python/Project.toml          |   4 +
 .../{python.jl => Python/runtests.jl}         |   4 +-
 test/nn/flux.jl                               |  25 -----
 test/runtests.jl                              | 100 +++++++++---------
 .../linear_algebra.jl                         |   0
 test/{integration => stdlibs}/random.jl       |   0
 53 files changed, 195 insertions(+), 107 deletions(-)
 rename test/{integration => common_integration}/special_functions.jl (100%)
 rename test/{ => core}/autodiff.jl (100%)
 rename test/{ => core}/basic.jl (100%)
 rename test/{ => core}/batching.jl (100%)
 rename test/{ => core}/bcast.jl (100%)
 rename test/{ => core}/buffer_donation.jl (100%)
 rename test/{ => core}/closure.jl (100%)
 rename test/{ => core}/compile.jl (100%)
 rename test/{ => core}/complex.jl (100%)
 rename test/{ => core}/config.jl (100%)
 rename test/{ => core}/constructor.jl (100%)
 rename test/{ => core}/control_flow.jl (100%)
 rename test/{ => core}/custom_number_types.jl (100%)
 rename test/{ => core}/indexing.jl (100%)
 rename test/{ => core}/ir.jl (100%)
 rename test/{ => core}/layout.jl (100%)
 rename test/{ => core}/ops.jl (100%)
 rename test/{ => core}/qa.jl (100%)
 rename test/{ => core}/sorting.jl (100%)
 rename test/{ => core}/struct.jl (100%)
 rename test/{ => core}/tracing.jl (100%)
 rename test/{ => core}/wrapped_arrays.jl (100%)
 rename test/{ => distributed}/cluster_detector.jl (100%)
 rename test/{ => distributed}/optimize_comm.jl (100%)
 rename test/{ => distributed}/sharding.jl (100%)
 create mode 100644 test/integration/CUDA/Project.toml
 rename test/integration/{ => CUDA}/cuda.jl (100%)
 rename test/integration/{ => CUDA}/kernelabstractions.jl (100%)
 create mode 100644 test/integration/CUDA/runtests.jl
 create mode 100644 test/integration/FFT/Project.toml
 rename test/integration/{fft.jl => FFT/runtests.jl} (100%)
 create mode 100644 test/integration/FancyArrayTypes/Project.toml
 rename test/integration/{ => FancyArrayTypes}/fillarrays.jl (100%)
 rename test/integration/{ => FancyArrayTypes}/offsetarrays.jl (88%)
 rename test/integration/{ => FancyArrayTypes}/onehotarrays.jl (100%)
 create mode 100644 test/integration/FancyArrayTypes/runtests.jl
 create mode 100644 test/integration/Flux/Project.toml
 create mode 100644 test/integration/Flux/runtests.jl
 create mode 100644 test/integration/Lux/Project.toml
 rename test/{nn => integration/Lux}/lux.jl (100%)
 rename test/{nn => integration/Lux}/luxlib.jl (100%)
 create mode 100644 test/integration/Lux/runtests.jl
 create mode 100644 test/integration/NNlib/Project.toml
 rename test/{nn/nnlib.jl => integration/NNlib/runtests.jl} (99%)
 create mode 100644 test/integration/Optimisers/Project.toml
 rename test/integration/{optimisers.jl => Optimisers/runtests.jl} (100%)
 create mode 100644 test/integration/Python/Project.toml
 rename test/integration/{python.jl => Python/runtests.jl} (98%)
 delete mode 100644 test/nn/flux.jl
 rename test/{integration => stdlibs}/linear_algebra.jl (100%)
 rename test/{integration => stdlibs}/random.jl (100%)

diff --git a/test/Project.toml b/test/Project.toml
index 05b1e8906a..efe76bc927 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -2,29 +2,17 @@
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
-FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Float8s = "81dfefd7-55b0-40c6-a251-db853704e186"
-Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 HypothesisTests = "09f84164-cd44-5f33-b23f-e6b0d136a0d5"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
-LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 MethodAnalysis = "85b6ec6f-f7df-4429-9514-a64bcd9ee824"
-NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
-OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
-Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
-PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -38,27 +26,15 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 Adapt = "4.1"
 ArrayInterface = "7.17.1"
-CUDA = "5.5"
 DLFP8Types = "0.1"
 Distributions = "0.25"
 Enzyme = "0.13.28"
-FFTW = "1.8"
-FillArrays = "1"
 Float8s = "0.1"
-Flux = "0.15, 0.16"
 Functors = "0.5"
 HypothesisTests = "0.11"
 InteractiveUtils = "1.10"
-KernelAbstractions = "0.9.30"
 LinearAlgebra = "1.10"
-Lux = "1.21"
-LuxLib = "1.11"
-NNlib = "0.9.26"
-OffsetArrays = "1"
-OneHotArrays = "0.2.6"
-Optimisers = "0.4"
 Preferences = "1.4"
-PythonCall = "0.9"
 Random = "1.10"
 Random123 = "1.7"
 SafeTestsets = "0.1"
diff --git a/test/integration/special_functions.jl b/test/common_integration/special_functions.jl
similarity index 100%
rename from test/integration/special_functions.jl
rename to test/common_integration/special_functions.jl
diff --git a/test/autodiff.jl b/test/core/autodiff.jl
similarity index 100%
rename from test/autodiff.jl
rename to test/core/autodiff.jl
diff --git a/test/basic.jl b/test/core/basic.jl
similarity index 100%
rename from test/basic.jl
rename to test/core/basic.jl
diff --git a/test/batching.jl b/test/core/batching.jl
similarity index 100%
rename from test/batching.jl
rename to test/core/batching.jl
diff --git a/test/bcast.jl b/test/core/bcast.jl
similarity index 100%
rename from test/bcast.jl
rename to test/core/bcast.jl
diff --git a/test/buffer_donation.jl b/test/core/buffer_donation.jl
similarity index 100%
rename from test/buffer_donation.jl
rename to test/core/buffer_donation.jl
diff --git a/test/closure.jl b/test/core/closure.jl
similarity index 100%
rename from test/closure.jl
rename to test/core/closure.jl
diff --git a/test/compile.jl b/test/core/compile.jl
similarity index 100%
rename from test/compile.jl
rename to test/core/compile.jl
diff --git a/test/complex.jl b/test/core/complex.jl
similarity index 100%
rename from test/complex.jl
rename to test/core/complex.jl
diff --git a/test/config.jl b/test/core/config.jl
similarity index 100%
rename from test/config.jl
rename to test/core/config.jl
diff --git a/test/constructor.jl b/test/core/constructor.jl
similarity index 100%
rename from test/constructor.jl
rename to test/core/constructor.jl
diff --git a/test/control_flow.jl b/test/core/control_flow.jl
similarity index 100%
rename from test/control_flow.jl
rename to test/core/control_flow.jl
diff --git a/test/custom_number_types.jl b/test/core/custom_number_types.jl
similarity index 100%
rename from test/custom_number_types.jl
rename to test/core/custom_number_types.jl
diff --git a/test/indexing.jl b/test/core/indexing.jl
similarity index 100%
rename from test/indexing.jl
rename to test/core/indexing.jl
diff --git a/test/ir.jl b/test/core/ir.jl
similarity index 100%
rename from test/ir.jl
rename to test/core/ir.jl
diff --git a/test/layout.jl b/test/core/layout.jl
similarity index 100%
rename from test/layout.jl
rename to test/core/layout.jl
diff --git a/test/ops.jl b/test/core/ops.jl
similarity index 100%
rename from test/ops.jl
rename to test/core/ops.jl
diff --git a/test/qa.jl b/test/core/qa.jl
similarity index 100%
rename from test/qa.jl
rename to test/core/qa.jl
diff --git a/test/sorting.jl b/test/core/sorting.jl
similarity index 100%
rename from test/sorting.jl
rename to test/core/sorting.jl
diff --git a/test/struct.jl b/test/core/struct.jl
similarity index 100%
rename from test/struct.jl
rename to test/core/struct.jl
diff --git a/test/tracing.jl b/test/core/tracing.jl
similarity index 100%
rename from test/tracing.jl
rename to test/core/tracing.jl
diff --git a/test/wrapped_arrays.jl b/test/core/wrapped_arrays.jl
similarity index 100%
rename from test/wrapped_arrays.jl
rename to test/core/wrapped_arrays.jl
diff --git a/test/cluster_detector.jl b/test/distributed/cluster_detector.jl
similarity index 100%
rename from test/cluster_detector.jl
rename to test/distributed/cluster_detector.jl
diff --git a/test/optimize_comm.jl b/test/distributed/optimize_comm.jl
similarity index 100%
rename from test/optimize_comm.jl
rename to test/distributed/optimize_comm.jl
diff --git a/test/sharding.jl b/test/distributed/sharding.jl
similarity index 100%
rename from test/sharding.jl
rename to test/distributed/sharding.jl
diff --git a/test/integration/CUDA/Project.toml b/test/integration/CUDA/Project.toml
new file mode 100644
index 0000000000..0c68429498
--- /dev/null
+++ b/test/integration/CUDA/Project.toml
@@ -0,0 +1,11 @@
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Reactant = {path = "../../.."}
+
+[extras]
+Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/cuda.jl b/test/integration/CUDA/cuda.jl
similarity index 100%
rename from test/integration/cuda.jl
rename to test/integration/CUDA/cuda.jl
diff --git a/test/integration/kernelabstractions.jl b/test/integration/CUDA/kernelabstractions.jl
similarity index 100%
rename from test/integration/kernelabstractions.jl
rename to test/integration/CUDA/kernelabstractions.jl
diff --git a/test/integration/CUDA/runtests.jl b/test/integration/CUDA/runtests.jl
new file mode 100644
index 0000000000..ee37d9f6b8
--- /dev/null
+++ b/test/integration/CUDA/runtests.jl
@@ -0,0 +1,11 @@
+using Test
+
+@testset "CUDA Integration" begin
+    @testset "CUDA" begin
+        include("cuda.jl")
+    end
+
+    @testset "KernelAbstractions" begin
+        include("kernelabstractions.jl")
+    end
+end
diff --git a/test/integration/FFT/Project.toml b/test/integration/FFT/Project.toml
new file mode 100644
index 0000000000..0b82a5346e
--- /dev/null
+++ b/test/integration/FFT/Project.toml
@@ -0,0 +1,10 @@
+[deps]
+FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Reactant = {path = "../../.."}
+
+[extras]
+Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/fft.jl b/test/integration/FFT/runtests.jl
similarity index 100%
rename from test/integration/fft.jl
rename to test/integration/FFT/runtests.jl
diff --git a/test/integration/FancyArrayTypes/Project.toml b/test/integration/FancyArrayTypes/Project.toml
new file mode 100644
index 0000000000..a258a55c27
--- /dev/null
+++ b/test/integration/FancyArrayTypes/Project.toml
@@ -0,0 +1,13 @@
+[deps]
+FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Reactant = {path = "../../.."}
+
+[extras]
+Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/fillarrays.jl b/test/integration/FancyArrayTypes/fillarrays.jl
similarity index 100%
rename from test/integration/fillarrays.jl
rename to test/integration/FancyArrayTypes/fillarrays.jl
diff --git a/test/integration/offsetarrays.jl b/test/integration/FancyArrayTypes/offsetarrays.jl
similarity index 88%
rename from test/integration/offsetarrays.jl
rename to test/integration/FancyArrayTypes/offsetarrays.jl
index a8bb4f899e..cf1d47a41f 100644
--- a/test/integration/offsetarrays.jl
+++ b/test/integration/FancyArrayTypes/offsetarrays.jl
@@ -1,10 +1,9 @@
-using Reactant
-using Test
-using OffsetArrays
+using Reactant, Test, OffsetArrays
 
 function scalar_index(x)
     @allowscalar getindex(x, -1, 0)
 end
+
 @testset "OffsetArrays" begin
     A = Float64.(reshape(1:15, 3, 5))
     OA = OffsetArray(A, -1:1, 0:4)
diff --git a/test/integration/onehotarrays.jl b/test/integration/FancyArrayTypes/onehotarrays.jl
similarity index 100%
rename from test/integration/onehotarrays.jl
rename to test/integration/FancyArrayTypes/onehotarrays.jl
diff --git a/test/integration/FancyArrayTypes/runtests.jl b/test/integration/FancyArrayTypes/runtests.jl
new file mode 100644
index 0000000000..d6d92bb013
--- /dev/null
+++ b/test/integration/FancyArrayTypes/runtests.jl
@@ -0,0 +1,15 @@
+using Test
+
+@testset "Fancy Array Types" begin
+    @testset "OffsetArrays" begin
+        include("offsetarrays.jl")
+    end
+
+    @testset "FillArrays" begin
+        include("fillarrays.jl")
+    end
+
+    @testset "OneHotArrays" begin
+        include("onehotarrays.jl")
+    end
+end
diff --git a/test/integration/Flux/Project.toml b/test/integration/Flux/Project.toml
new file mode 100644
index 0000000000..d0a6f6687d
--- /dev/null
+++ b/test/integration/Flux/Project.toml
@@ -0,0 +1,10 @@
+[deps]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Reactant = {path = "../../.."}
+
+[extras]
+Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/Flux/runtests.jl b/test/integration/Flux/runtests.jl
new file mode 100644
index 0000000000..99ef9b67f3
--- /dev/null
+++ b/test/integration/Flux/runtests.jl
@@ -0,0 +1,24 @@
+using Reactant, Flux, Test
+
+@testset "Flux.jl Integration" begin
+    noisy = rand(Float32, 2, 1000)
+    truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]
+
+    model = Chain(
+        Dense(2 => 3, tanh),
+        BatchNorm(3),
+        Dense(3 => 2),
+        softmax,
+    )
+
+    origout = model(noisy)
+
+    cmodel = Reactant.to_rarray(model)
+    cnoisy = Reactant.to_rarray(noisy)
+
+    fn = (a, b) -> a(b)
+    f = @compile fn(cmodel, cnoisy)
+
+    comp = f(cmodel, cnoisy)
+    @test origout ≈ comp atol = 1e-3 rtol = 1e-2
+end
diff --git a/test/integration/Lux/Project.toml b/test/integration/Lux/Project.toml
new file mode 100644
index 0000000000..d08afcddbe
--- /dev/null
+++ b/test/integration/Lux/Project.toml
@@ -0,0 +1,17 @@
+[deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Reactant = {path = "../../.."}
+
+[extras]
+Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/nn/lux.jl b/test/integration/Lux/lux.jl
similarity index 100%
rename from test/nn/lux.jl
rename to test/integration/Lux/lux.jl
diff --git a/test/nn/luxlib.jl b/test/integration/Lux/luxlib.jl
similarity index 100%
rename from test/nn/luxlib.jl
rename to test/integration/Lux/luxlib.jl
diff --git a/test/integration/Lux/runtests.jl b/test/integration/Lux/runtests.jl
new file mode 100644
index 0000000000..027fa5ad04
--- /dev/null
+++ b/test/integration/Lux/runtests.jl
@@ -0,0 +1,11 @@
+using Test
+
+@testset "Lux.jl Integration" begin
+    @testset "LuxLib Primitives" begin
+        include("luxlib.jl")
+    end
+
+    @testset "Lux Integration" begin
+        include("lux.jl")
+    end
+end
diff --git a/test/integration/NNlib/Project.toml b/test/integration/NNlib/Project.toml
new file mode 100644
index 0000000000..cb00460300
--- /dev/null
+++ b/test/integration/NNlib/Project.toml
@@ -0,0 +1,11 @@
+[deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+
+[sources]
+Reactant = {path = "../../.."}
+
+[extras]
+Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/nn/nnlib.jl b/test/integration/NNlib/runtests.jl
similarity index 99%
rename from test/nn/nnlib.jl
rename to test/integration/NNlib/runtests.jl
index 240c52c155..876ee3eb9d 100644
--- a/test/nn/nnlib.jl
+++ b/test/integration/NNlib/runtests.jl
@@ -1,5 +1,4 @@
-using NNlib, Reactant, Enzyme
-using Statistics
+using NNlib, Reactant, Enzyme, Test, Statistics
 
 @testset "Activation Functions" begin
     sumabs2(f, x) = sum(abs2, f.(x))
diff --git a/test/integration/Optimisers/Project.toml b/test/integration/Optimisers/Project.toml
new file mode 100644
index 0000000000..8f338853cd
--- /dev/null
+++ b/test/integration/Optimisers/Project.toml
@@ -0,0 +1,4 @@
+[deps]
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/integration/optimisers.jl b/test/integration/Optimisers/runtests.jl
similarity index 100%
rename from test/integration/optimisers.jl
rename to test/integration/Optimisers/runtests.jl
diff --git a/test/integration/Python/Project.toml b/test/integration/Python/Project.toml
new file mode 100644
index 0000000000..01c941ed47
--- /dev/null
+++ b/test/integration/Python/Project.toml
@@ -0,0 +1,4 @@
+[deps]
+PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
+Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/integration/python.jl b/test/integration/Python/runtests.jl
similarity index 98%
rename from test/integration/python.jl
rename to test/integration/Python/runtests.jl
index f0f04ae7c8..c51eb6e7b3 100644
--- a/test/integration/python.jl
+++ b/test/integration/Python/runtests.jl
@@ -1,8 +1,6 @@
-using Reactant
+using Reactant, Test
 using Reactant: Ops
 
-using Test
-
 fn(x, y) = sin.(x) .+ cos.(y.x[1:2, :])
 
 # Jax on Github CI dislikes X86 macos
diff --git a/test/nn/flux.jl b/test/nn/flux.jl
deleted file mode 100644
index f4c314bccd..0000000000
--- a/test/nn/flux.jl
+++ /dev/null
@@ -1,25 +0,0 @@
-using Reactant, Flux
-
-@testset "Flux.jl Integration" begin
-    # Generate some data for the XOR problem: vectors of length 2, as columns of a matrix:
-    noisy = rand(Float32, 2, 1000)                                        # 2×1000 Matrix{Float32}
-    truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]   # 1000-element Vector{Bool}
-
-    # Define our model, a multi-layer perceptron with one hidden layer of size 3:
-    model = Chain(
-        Dense(2 => 3, tanh),   # activation function inside layer
-        BatchNorm(3),
-        Dense(3 => 2),
-        softmax,
-    )
-
-    origout = model(noisy)
-
-    cmodel = Reactant.to_rarray(model)
-    cnoisy = Reactant.to_rarray(noisy)
-
-    f = Reactant.compile((a, b) -> a(b), (cmodel, cnoisy))
-
-    comp = f(cmodel, cnoisy)
-    @test origout ≈ comp atol = 1e-3 rtol = 1e-2
-end
diff --git a/test/runtests.jl b/test/runtests.jl
index 98a02a7de0..fc3a672171 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,66 +1,66 @@
 using Reactant, SafeTestsets, Test
 
-if lowercase(get(ENV, "REACTANT_BACKEND_GROUP", "all")) == "gpu"
+const REACTANT_BACKEND_GROUP = lowercase(get(ENV, "REACTANT_BACKEND_GROUP", "all"))
+
+if REACTANT_BACKEND_GROUP == "gpu"
     Reactant.set_default_backend("gpu")
+elseif REACTANT_BACKEND_GROUP == "tpu"
+    Reactant.set_default_backend("tpu")
+elseif REACTANT_BACKEND_GROUP == "cpu"
+    Reactant.set_default_backend("cpu")
 end
 
-const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
-
 @testset "Reactant.jl Tests" begin
-    if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "core"
-        if Sys.isapple() && haskey(Reactant.XLA.global_backend_state.clients, "metal")
-            @safetestset "Metal Plugin" include("plugins/metal.jl")
-        end
+    @testset "Core" begin
+        @safetestset "Layout" include("core/layout.jl")
+        @safetestset "Tracing" include("core/tracing.jl")
+        @safetestset "Basic" include("core/basic.jl")
+        @safetestset "Constructor" include("core/constructor.jl")
+        @safetestset "Autodiff" include("core/autodiff.jl")
+        @safetestset "Complex" include("core/complex.jl")
+        @safetestset "Broadcast" include("core/bcast.jl")
+        @safetestset "Struct" include("core/struct.jl")
+        @safetestset "Closure" include("core/closure.jl")
+        @safetestset "Compile" include("core/compile.jl")
+        @safetestset "IR" include("core/ir.jl")
+        @safetestset "Buffer Donation" include("core/buffer_donation.jl")
+        @safetestset "Wrapped Arrays" include("core/wrapped_arrays.jl")
+        @safetestset "Control Flow" include("core/control_flow.jl")
+        @safetestset "Sorting" include("core/sorting.jl")
+        @safetestset "Shortcuts to MLIR ops" include("core/ops.jl")
+        @safetestset "Indexing" include("core/indexing.jl")
+        @safetestset "Config" include("core/config.jl")
+        @safetestset "Batching" include("core/batching.jl")
 
-        @safetestset "Layout" include("layout.jl")
-        @safetestset "Tracing" include("tracing.jl")
-        @safetestset "Basic" include("basic.jl")
-        @safetestset "Constructor" include("constructor.jl")
-        @safetestset "Autodiff" include("autodiff.jl")
-        @safetestset "Complex" include("complex.jl")
-        @safetestset "Broadcast" include("bcast.jl")
-        @safetestset "Struct" include("struct.jl")
-        @safetestset "Closure" include("closure.jl")
-        @safetestset "Compile" include("compile.jl")
-        @safetestset "IR" include("ir.jl")
-        @safetestset "Buffer Donation" include("buffer_donation.jl")
-        @safetestset "Wrapped Arrays" include("wrapped_arrays.jl")
-        @safetestset "Control Flow" include("control_flow.jl")
-        @safetestset "Sorting" include("sorting.jl")
-        @safetestset "Shortcuts to MLIR ops" include("ops.jl")
-        @safetestset "Indexing" include("indexing.jl")
         if !Sys.isapple()
-            @safetestset "Custom Number Types" include("custom_number_types.jl")
+            @safetestset "Custom Number Types" include("core/custom_number_types.jl")
         end
-        @safetestset "Sharding" include("sharding.jl")
-        @safetestset "Comm Optimization" include("optimize_comm.jl")
-        @safetestset "Cluster Detection" include("cluster_detector.jl")
-        @safetestset "Config" include("config.jl")
-        @safetestset "Batching" include("batching.jl")
-        @safetestset "QA" include("qa.jl")
     end
 
-    if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"
-        @safetestset "CUDA" include("integration/cuda.jl")
-        @safetestset "KernelAbstractions" include("integration/kernelabstractions.jl")
-        @safetestset "Linear Algebra" include("integration/linear_algebra.jl")
-        @safetestset "OffsetArrays" include("integration/offsetarrays.jl")
-        @safetestset "OneHotArrays" include("integration/onehotarrays.jl")
-        @safetestset "AbstractFFTs" include("integration/fft.jl")
-        @safetestset "SpecialFunctions" include("integration/special_functions.jl")
-        @safetestset "Random" include("integration/random.jl")
-        @safetestset "Python" include("integration/python.jl")
-        @safetestset "Optimisers" include("integration/optimisers.jl")
-        @safetestset "FillArrays" include("integration/fillarrays.jl")
-        @safetestset "Zygote" include("integration/zygote.jl")
+    @testset "Distributed" begin
+        @safetestset "Sharding" include("distributed/sharding.jl")
+        @safetestset "Comm Optimization" include("distributed/optimize_comm.jl")
+        @safetestset "Cluster Detection" include("distributed/cluster_detector.jl")
     end
 
-    if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "neural_networks"
-        @safetestset "NNlib Primitives" include("nn/nnlib.jl")
-        @safetestset "Flux.jl Integration" include("nn/flux.jl")
-        if Sys.islinux()
-            @safetestset "LuxLib Primitives" include("nn/luxlib.jl")
-            @safetestset "Lux Integration" include("nn/lux.jl")
+    @testset "Plugins" begin
+        if Sys.isapple() && haskey(Reactant.XLA.global_backend_state.clients, "metal")
+            @safetestset "Metal Plugin" include("plugins/metal.jl")
         end
     end
+
+    @testset "Standard Libraries" begin
+        @safetestset "Linear Algebra" include("stdlibs/linear_algebra.jl")
+        @safetestset "Random" include("stdlibs/random.jl")
+    end
+
+    @testset "Common Integrations" begin
+        # most integrations are tested in the integration tests except deps that are
+        # very common
+        @safetestset "SpecialFunctions" include("common_integration/special_functions.jl")
+    end
+
+    @testset "Quality Assurance" begin
+        @safetestset "QA" include("core/qa.jl")
+    end
 end
diff --git a/test/integration/linear_algebra.jl b/test/stdlibs/linear_algebra.jl
similarity index 100%
rename from test/integration/linear_algebra.jl
rename to test/stdlibs/linear_algebra.jl
diff --git a/test/integration/random.jl b/test/stdlibs/random.jl
similarity index 100%
rename from test/integration/random.jl
rename to test/stdlibs/random.jl

From 730eb1b926b69b78e43b444376237c8334bd201f Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 13:13:48 -0400
Subject: [PATCH 05/16] test: remove test groups

---
 .github/workflows/CI.yml                      | 15 -----------
 test/Project.toml                             |  2 --
 test/integration/Flux/Project.toml            |  1 +
 test/integration/Flux/flux.jl                 | 24 +++++++++++++++++
 test/integration/Flux/runtests.jl             | 27 +++++--------------
 test/integration/{ => Flux}/zygote.jl         |  0
 test/integration/Lux/Project.toml             |  1 +
 .../runtests.jl => Lux/optimisers.jl}         |  0
 test/integration/Lux/runtests.jl              |  4 +++
 test/integration/Optimisers/Project.toml      |  4 ---
 10 files changed, 37 insertions(+), 41 deletions(-)
 create mode 100644 test/integration/Flux/flux.jl
 rename test/integration/{ => Flux}/zygote.jl (100%)
 rename test/integration/{Optimisers/runtests.jl => Lux/optimisers.jl} (100%)
 delete mode 100644 test/integration/Optimisers/Project.toml

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index b60f801cbc..ae375e4654 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -48,10 +48,6 @@ jobs:
           - macOS-latest
           - windows-latest
           - linux-x86-ct6e-180-4tpu
-        test_group:
-          - core
-          - neural_networks
-          - integration
         runtime:
           - "pjrt"
           - "ifrt"
@@ -65,7 +61,6 @@ jobs:
       julia_version: ${{ matrix.version }}
       os: ${{ matrix.os }}
       runtime: ${{ matrix.runtime }}
-      test_group: ${{ matrix.test_group }}
 
   # This has been broken for a while, originating from CUDA.jl
   # test-assertions:
@@ -74,26 +69,17 @@ jobs:
   #     matrix:
   #       version:
   #         - "1.10"
-  #       test_group:
-  #         - core
-  #         - neural_networks
-  #         - integration
   #   uses: ./.github/workflows/CommonCI.yml
   #   with:
   #     julia_version: ${{ matrix.version }}
   #     os: "ubuntu-24.04"
   #     runtime: "both"
   #     assertions: true
-  #     test_group: ${{ matrix.test_group }}
 
   downgrade:
     strategy:
       fail-fast: false
       matrix:
-        test_group:
-          - core
-          - neural_networks
-          - integration
         runtime:
           - "pjrt"
           - "ifrt"
@@ -102,5 +88,4 @@ jobs:
       julia_version: "1.10"
       os: "ubuntu-24.04"
       runtime: ${{ matrix.runtime }}
-      test_group: ${{ matrix.test_group }}
       downgrade_testing: true
diff --git a/test/Project.toml b/test/Project.toml
index efe76bc927..8bd2903feb 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -21,7 +21,6 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 Adapt = "4.1"
@@ -43,7 +42,6 @@ StableRNGs = "1"
 Statistics = "1.10"
 StatsBase = "0.34"
 Test = "1.10"
-Zygote = "0.7"
 
 [extras]
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/Flux/Project.toml b/test/integration/Flux/Project.toml
index d0a6f6687d..db3c32f98a 100644
--- a/test/integration/Flux/Project.toml
+++ b/test/integration/Flux/Project.toml
@@ -2,6 +2,7 @@
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [sources]
 Reactant = {path = "../../.."}
diff --git a/test/integration/Flux/flux.jl b/test/integration/Flux/flux.jl
new file mode 100644
index 0000000000..99ef9b67f3
--- /dev/null
+++ b/test/integration/Flux/flux.jl
@@ -0,0 +1,24 @@
+using Reactant, Flux, Test
+
+@testset "Flux.jl Integration" begin
+    noisy = rand(Float32, 2, 1000)
+    truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]
+
+    model = Chain(
+        Dense(2 => 3, tanh),
+        BatchNorm(3),
+        Dense(3 => 2),
+        softmax,
+    )
+
+    origout = model(noisy)
+
+    cmodel = Reactant.to_rarray(model)
+    cnoisy = Reactant.to_rarray(noisy)
+
+    fn = (a, b) -> a(b)
+    f = @compile fn(cmodel, cnoisy)
+
+    comp = f(cmodel, cnoisy)
+    @test origout ≈ comp atol = 1e-3 rtol = 1e-2
+end
diff --git a/test/integration/Flux/runtests.jl b/test/integration/Flux/runtests.jl
index 99ef9b67f3..8c05868052 100644
--- a/test/integration/Flux/runtests.jl
+++ b/test/integration/Flux/runtests.jl
@@ -1,24 +1,11 @@
-using Reactant, Flux, Test
+using Test
 
 @testset "Flux.jl Integration" begin
-    noisy = rand(Float32, 2, 1000)
-    truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]
+    @testset "Zygote Integration" begin
+        include("zygote.jl")
+    end
 
-    model = Chain(
-        Dense(2 => 3, tanh),
-        BatchNorm(3),
-        Dense(3 => 2),
-        softmax,
-    )
-
-    origout = model(noisy)
-
-    cmodel = Reactant.to_rarray(model)
-    cnoisy = Reactant.to_rarray(noisy)
-
-    fn = (a, b) -> a(b)
-    f = @compile fn(cmodel, cnoisy)
-
-    comp = f(cmodel, cnoisy)
-    @test origout ≈ comp atol = 1e-3 rtol = 1e-2
+    @testset "Flux Integration" begin
+        include("flux.jl")
+    end
 end
diff --git a/test/integration/zygote.jl b/test/integration/Flux/zygote.jl
similarity index 100%
rename from test/integration/zygote.jl
rename to test/integration/Flux/zygote.jl
diff --git a/test/integration/Lux/Project.toml b/test/integration/Lux/Project.toml
index d08afcddbe..7bef50ba53 100644
--- a/test/integration/Lux/Project.toml
+++ b/test/integration/Lux/Project.toml
@@ -5,6 +5,7 @@ Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
diff --git a/test/integration/Optimisers/runtests.jl b/test/integration/Lux/optimisers.jl
similarity index 100%
rename from test/integration/Optimisers/runtests.jl
rename to test/integration/Lux/optimisers.jl
diff --git a/test/integration/Lux/runtests.jl b/test/integration/Lux/runtests.jl
index 027fa5ad04..173f2450ab 100644
--- a/test/integration/Lux/runtests.jl
+++ b/test/integration/Lux/runtests.jl
@@ -8,4 +8,8 @@ using Test
     @testset "Lux Integration" begin
         include("lux.jl")
     end
+
+    @testset "Optimisers Integration" begin
+        include("optimisers.jl")
+    end
 end
diff --git a/test/integration/Optimisers/Project.toml b/test/integration/Optimisers/Project.toml
deleted file mode 100644
index 8f338853cd..0000000000
--- a/test/integration/Optimisers/Project.toml
+++ /dev/null
@@ -1,4 +0,0 @@
-[deps]
-Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
-Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

From 5670ede9689d6e3cd3f1228695468423cf1b2d26 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 13:30:31 -0400
Subject: [PATCH 06/16] ci: setup integration testing

---
 .github/workflows/CI.yml                      | 20 +++----
 .github/workflows/CommonCI.yml                | 36 +++++++++++-
 .github/workflows/Integration.yml             | 55 +++++++++++++++++++
 test/integration/CUDA/Project.toml            |  3 -
 test/integration/FFT/Project.toml             |  3 -
 test/integration/FancyArrayTypes/Project.toml |  3 -
 test/integration/Flux/Project.toml            |  3 -
 test/integration/Lux/Project.toml             |  3 -
 test/integration/NNlib/Project.toml           |  3 -
 test/integration/Python/Project.toml          |  3 +
 10 files changed, 98 insertions(+), 34 deletions(-)
 create mode 100644 .github/workflows/Integration.yml

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index ae375e4654..546eadfecd 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -62,19 +62,13 @@ jobs:
       os: ${{ matrix.os }}
       runtime: ${{ matrix.runtime }}
 
-  # This has been broken for a while, originating from CUDA.jl
-  # test-assertions:
-  #   strategy:
-  #     fail-fast: false
-  #     matrix:
-  #       version:
-  #         - "1.10"
-  #   uses: ./.github/workflows/CommonCI.yml
-  #   with:
-  #     julia_version: ${{ matrix.version }}
-  #     os: "ubuntu-24.04"
-  #     runtime: "both"
-  #     assertions: true
+  test-assertions:
+    uses: ./.github/workflows/CommonCI.yml
+    with:
+      julia_version: "1.10"
+      os: "ubuntu-24.04"
+      runtime: "both"
+      assertions: true
 
   downgrade:
     strategy:
diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml
index 5bb654bf5d..84bc6869e3 100644
--- a/.github/workflows/CommonCI.yml
+++ b/.github/workflows/CommonCI.yml
@@ -35,6 +35,11 @@ on:
         required: false
         default: false
         type: boolean
+      project:
+        description: "Project"
+        required: false
+        default: "."
+        type: string
 
 jobs:
   test:
@@ -132,9 +137,22 @@ jobs:
           end
           Pkg.develop(dev_pks)
         shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
-        # Only in Julia v1.10 we need to install `ReactantCore` manually.
         if: ${{ inputs.julia_version == '1.10' || inputs.julia_version == 'lts' }}
 
+      - name: Install Project
+        if: ${{ inputs.project != '.' }}
+        run: |
+          import Pkg
+          PKg.activate(get(ENV, "JULIA_PROJECT", nothing))
+          if VERSION < v"1.11-"
+              Pkg.develop([PackageSpec(; path="."), PackageSpec(; path="lib/ReactantCore")])
+          end
+          Pkg.instantiate()
+        shell: julia --color=yes --code-coverage=user --depwarn=yes {0}
+        env:
+          JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
+          JULIA_PROJECT: ${{ inputs.project }}
+
       # Run the tests (pjrt or ifrt or both)
       - name: "Setup Runtime Preferences (PJRT)"
         if: ${{ inputs.runtime == 'pjrt' || inputs.runtime == 'both' }}
@@ -145,8 +163,9 @@ jobs:
           contents: |
             [Reactant]
             xla_runtime = "PJRT"
+
       - name: "Run Tests (PJRT)"
-        if: ${{ inputs.runtime == 'pjrt' || inputs.runtime == 'both' }}
+        if: ${{ (inputs.runtime == 'pjrt' || inputs.runtime == 'both') && inputs.project == '.' }}
         timeout-minutes: 120
         run: |
           import Pkg
@@ -159,6 +178,11 @@ jobs:
         env:
           ALLOW_RERESOLVE: ${{ !inputs.downgrade_testing }}
           REACTANT_TEST_GROUP: ${{ inputs.test_group }}
+      - name: "Run Integration Tests (PJRT)"
+        if: ${{ (inputs.runtime == 'pjrt' || inputs.runtime == 'both') && inputs.project != '.' }}
+        timeout-minutes: 30
+        run: julia --color=yes --project=test/integration/${{ inputs.project }} --threads=auto --check-bounds=yes -O1 test/integration/${{ inputs.project }}/runtests.jl
+        shell: bash
 
       - name: "Setup Runtime Preferences (IFRT)"
         if: ${{ inputs.runtime == 'ifrt' || inputs.runtime == 'both' }}
@@ -169,8 +193,9 @@ jobs:
           contents: |
             [Reactant]
             xla_runtime = "IFRT"
+
       - name: "Run Tests (IFRT)"
-        if: ${{ inputs.runtime == 'ifrt' || inputs.runtime == 'both' }}
+        if: ${{ (inputs.runtime == 'ifrt' || inputs.runtime == 'both') && inputs.project == '.' }}
         timeout-minutes: 120
         run: |
           import Pkg
@@ -183,6 +208,11 @@ jobs:
         env:
           ALLOW_RERESOLVE: ${{ !inputs.downgrade_testing }}
           REACTANT_TEST_GROUP: ${{ inputs.test_group }}
+      - name: "Run Integration Tests (IFRT)"
+        if: ${{ (inputs.runtime == 'ifrt' || inputs.runtime == 'both') && inputs.project != '.' }}
+        timeout-minutes: 30
+        run: julia --color=yes --project=test/integration/${{ inputs.project }} --threads=auto --check-bounds=yes -O1 test/integration/${{ inputs.project }}/runtests.jl
+        shell: bash
 
       - name: "Upload MLIR modules"
         uses: actions/upload-artifact@v4
diff --git a/.github/workflows/Integration.yml b/.github/workflows/Integration.yml
new file mode 100644
index 0000000000..7bb970182a
--- /dev/null
+++ b/.github/workflows/Integration.yml
@@ -0,0 +1,55 @@
+name: Integration
+
+on:
+  pull_request:
+    paths:
+      - ".github/workflows/Integration.yml"
+      - "ext/**"
+      - "lib/**"
+      - "src/**"
+      - "test/**"
+      - "Project.toml"
+  push:
+    branches:
+      - main
+      - release-*
+    tags: ["*"]
+    paths:
+      - ".github/workflows/Integration.yml"
+      - "ext/**"
+      - "lib/**"
+      - "src/**"
+      - "test/**"
+      - "Project.toml"
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  integration:
+    strategy:
+      fail-fast: false
+      matrix:
+        project:
+          - CUDA
+          - FancyArrayTypes
+          - FFT
+          - Flux
+          - Lux
+          - NNlib
+          - Python
+        version:
+          - "1.10"
+          - "1.11"
+        os:
+          - ubuntu-latest
+          - linux-x86-ct6e-180-4tpu
+    uses: ./.github/workflows/CommonCI.yml
+    with:
+      julia_version: ${{ matrix.version }}
+      os: ${{ matrix.os }}
+      runtime: "both"
+      project: "test/integration/${{ matrix.project }}"
diff --git a/test/integration/CUDA/Project.toml b/test/integration/CUDA/Project.toml
index 0c68429498..f9363a2916 100644
--- a/test/integration/CUDA/Project.toml
+++ b/test/integration/CUDA/Project.toml
@@ -6,6 +6,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 Reactant = {path = "../../.."}
-
-[extras]
-Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/FFT/Project.toml b/test/integration/FFT/Project.toml
index 0b82a5346e..a93b2fc311 100644
--- a/test/integration/FFT/Project.toml
+++ b/test/integration/FFT/Project.toml
@@ -5,6 +5,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 Reactant = {path = "../../.."}
-
-[extras]
-Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/FancyArrayTypes/Project.toml b/test/integration/FancyArrayTypes/Project.toml
index a258a55c27..dde4fb4800 100644
--- a/test/integration/FancyArrayTypes/Project.toml
+++ b/test/integration/FancyArrayTypes/Project.toml
@@ -8,6 +8,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 Reactant = {path = "../../.."}
-
-[extras]
-Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/Flux/Project.toml b/test/integration/Flux/Project.toml
index db3c32f98a..28dd0fdfea 100644
--- a/test/integration/Flux/Project.toml
+++ b/test/integration/Flux/Project.toml
@@ -6,6 +6,3 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [sources]
 Reactant = {path = "../../.."}
-
-[extras]
-Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/Lux/Project.toml b/test/integration/Lux/Project.toml
index 7bef50ba53..0f85c564e2 100644
--- a/test/integration/Lux/Project.toml
+++ b/test/integration/Lux/Project.toml
@@ -13,6 +13,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 Reactant = {path = "../../.."}
-
-[extras]
-Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/NNlib/Project.toml b/test/integration/NNlib/Project.toml
index cb00460300..b4b1da67c2 100644
--- a/test/integration/NNlib/Project.toml
+++ b/test/integration/NNlib/Project.toml
@@ -6,6 +6,3 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [sources]
 Reactant = {path = "../../.."}
-
-[extras]
-Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
diff --git a/test/integration/Python/Project.toml b/test/integration/Python/Project.toml
index 01c941ed47..a24e261e4b 100644
--- a/test/integration/Python/Project.toml
+++ b/test/integration/Python/Project.toml
@@ -2,3 +2,6 @@
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[sources]
+Reactant = {path = "../../.."}

From 7364a30fcb72d9fbee7b1b58832295f49a862d03 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 13:30:50 -0400
Subject: [PATCH 07/16] ci: disable expensive runners while testing (drop me)

---
 .github/workflows/Integration.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/Integration.yml b/.github/workflows/Integration.yml
index 7bb970182a..fb90d01c76 100644
--- a/.github/workflows/Integration.yml
+++ b/.github/workflows/Integration.yml
@@ -46,7 +46,8 @@ jobs:
           - "1.11"
         os:
           - ubuntu-latest
-          - linux-x86-ct6e-180-4tpu
+        # TODO: enable these tests
+        #   - linux-x86-ct6e-180-4tpu
     uses: ./.github/workflows/CommonCI.yml
     with:
       julia_version: ${{ matrix.version }}

From 83729a60aa77446509b90fe1b13bd039dd0d08f9 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 13:32:29 -0400
Subject: [PATCH 08/16] fix: typo

---
 .github/workflows/CommonCI.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml
index 84bc6869e3..4a3f0be693 100644
--- a/.github/workflows/CommonCI.yml
+++ b/.github/workflows/CommonCI.yml
@@ -143,7 +143,7 @@ jobs:
         if: ${{ inputs.project != '.' }}
         run: |
           import Pkg
-          PKg.activate(get(ENV, "JULIA_PROJECT", nothing))
+          Pkg.activate(get(ENV, "JULIA_PROJECT", nothing))
           if VERSION < v"1.11-"
               Pkg.develop([PackageSpec(; path="."), PackageSpec(; path="lib/ReactantCore")])
           end

From e186cfb1d0cd4958efebd1d01a3cf7395216ef5b Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 13:32:51 -0400
Subject: [PATCH 09/16] chore: run formatter

---
 .github/workflows/CommonCI.yml | 5 ++++-
 test/integration/Flux/flux.jl  | 7 +------
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml
index 4a3f0be693..c623d86cbd 100644
--- a/.github/workflows/CommonCI.yml
+++ b/.github/workflows/CommonCI.yml
@@ -145,7 +145,10 @@ jobs:
           import Pkg
           Pkg.activate(get(ENV, "JULIA_PROJECT", nothing))
           if VERSION < v"1.11-"
-              Pkg.develop([PackageSpec(; path="."), PackageSpec(; path="lib/ReactantCore")])
+              Pkg.develop([
+                Pkg.PackageSpec(; path="."),
+                Pkg.PackageSpec(; path="lib/ReactantCore")
+              ])
           end
           Pkg.instantiate()
         shell: julia --color=yes --code-coverage=user --depwarn=yes {0}
diff --git a/test/integration/Flux/flux.jl b/test/integration/Flux/flux.jl
index 99ef9b67f3..b94b5946f2 100644
--- a/test/integration/Flux/flux.jl
+++ b/test/integration/Flux/flux.jl
@@ -4,12 +4,7 @@ using Reactant, Flux, Test
     noisy = rand(Float32, 2, 1000)
     truth = [xor(col[1] > 0.5, col[2] > 0.5) for col in eachcol(noisy)]
 
-    model = Chain(
-        Dense(2 => 3, tanh),
-        BatchNorm(3),
-        Dense(3 => 2),
-        softmax,
-    )
+    model = Chain(Dense(2 => 3, tanh), BatchNorm(3), Dense(3 => 2), softmax)
 
     origout = model(noisy)
 

From 0f21e25b4041b48aca380de0cd692d981124e6d1 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 13:40:53 -0400
Subject: [PATCH 10/16] ci: fix

---
 .buildkite/pipeline.yml            |  7 +---
 .github/workflows/CI.yml           | 47 +++++++++++++++++++++++++
 .github/workflows/CommonCI.yml     |  6 ++--
 .github/workflows/Integration.yml  | 56 ------------------------------
 test/integration/Flux/Project.toml |  1 +
 5 files changed, 52 insertions(+), 65 deletions(-)
 delete mode 100644 .github/workflows/Integration.yml

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 44f8b50dd7..7ba121da05 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -1,15 +1,11 @@
 steps:
   - group: ":test_tube: Tests"
     steps:
-      - label: ":julia: :linux: CUDA Julia v{{matrix.version}} -- {{matrix.group}} -- {{matrix.runtime}}"
+      - label: ":julia: :linux: CUDA Julia v{{matrix.version}} -- {{matrix.runtime}}"
         matrix:
           setup:
             version:
               - "1.10"
-            group:
-              - core
-              - neural_networks
-              - integration
             runtime:
               - "PJRT"
               - "IFRT"
@@ -41,7 +37,6 @@ steps:
           queue: "juliagpu"
           cuda: "*"
         env:
-          REACTANT_TEST_GROUP: "{{matrix.group}}"
           JULIA_DEBUG: "Reactant,Reactant_jll"
           CUDA_VISIBLE_DEVICES: 0
           REACTANT_BACKEND_GROUP: "GPU"
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 546eadfecd..bd86ebb2eb 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -83,3 +83,50 @@ jobs:
       os: "ubuntu-24.04"
       runtime: ${{ matrix.runtime }}
       downgrade_testing: true
+
+  integration-cpu:
+    strategy:
+      fail-fast: false
+      matrix:
+        project:
+          - FancyArrayTypes
+          - FFT
+          - Python
+        version:
+          - "1.10"
+          - "1.11"
+    uses: ./.github/workflows/CommonCI.yml
+    with:
+      julia_version: ${{ matrix.version }}
+      os: "ubuntu-latest"
+      runtime: "both"
+      project: "test/integration/${{ matrix.project }}"
+
+  integration:
+    strategy:
+      fail-fast: false
+      matrix:
+        project:
+          - CUDA
+          - Flux
+          - Lux
+          - NNlib
+        version:
+          - "1.10"
+          - "1.11"
+        os:
+          - ubuntu-latest
+          - linux-x86-ct6e-180-4tpu
+        include:
+          - project: CUDA
+            os: linux-x86-a2-48-a100-4gpu
+            version: "1.10"
+          - project: CUDA
+            os: linux-x86-a2-48-a100-4gpu
+            version: "1.11"
+    uses: ./.github/workflows/CommonCI.yml
+    with:
+      julia_version: ${{ matrix.version }}
+      os: ${{ matrix.os }}
+      runtime: "both"
+      project: "test/integration/${{ matrix.project }}"
diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml
index c623d86cbd..80484e8c3c 100644
--- a/.github/workflows/CommonCI.yml
+++ b/.github/workflows/CommonCI.yml
@@ -75,7 +75,7 @@ jobs:
       - uses: julia-actions/cache@v2
         id: julia-cache
         with:
-          cache-name: julia-cache;workflow=${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.runtime }}-assertions=${{ inputs.assertions }}-${{ github.event_name }}-${{ inputs.test_group }}-${{ inputs.downgrade_testing }}-${{ inputs.localjll }}
+          cache-name: julia-cache;workflow=${{ inputs.julia_version }}-${{ inputs.os }}-${{ inputs.runtime }}-assertions=${{ inputs.assertions }}-${{ github.event_name }}-${{ inputs.test_group }}-${{ inputs.downgrade_testing }}-${{ inputs.localjll }}-${{ inputs.project }}
 
       - uses: julia-actions/julia-downgrade-compat@v2
         if: ${{ inputs.downgrade_testing }}
@@ -184,7 +184,7 @@ jobs:
       - name: "Run Integration Tests (PJRT)"
         if: ${{ (inputs.runtime == 'pjrt' || inputs.runtime == 'both') && inputs.project != '.' }}
         timeout-minutes: 30
-        run: julia --color=yes --project=test/integration/${{ inputs.project }} --threads=auto --check-bounds=yes -O1 test/integration/${{ inputs.project }}/runtests.jl
+        run: julia --color=yes --project=${{ inputs.project }} --threads=auto --check-bounds=yes -O1 ${{ inputs.project }}/runtests.jl
         shell: bash
 
       - name: "Setup Runtime Preferences (IFRT)"
@@ -214,7 +214,7 @@ jobs:
       - name: "Run Integration Tests (IFRT)"
         if: ${{ (inputs.runtime == 'ifrt' || inputs.runtime == 'both') && inputs.project != '.' }}
         timeout-minutes: 30
-        run: julia --color=yes --project=test/integration/${{ inputs.project }} --threads=auto --check-bounds=yes -O1 test/integration/${{ inputs.project }}/runtests.jl
+        run: julia --color=yes --project=${{ inputs.project }} --threads=auto --check-bounds=yes -O1 ${{ inputs.project }}/runtests.jl
         shell: bash
 
       - name: "Upload MLIR modules"
diff --git a/.github/workflows/Integration.yml b/.github/workflows/Integration.yml
deleted file mode 100644
index fb90d01c76..0000000000
--- a/.github/workflows/Integration.yml
+++ /dev/null
@@ -1,56 +0,0 @@
-name: Integration
-
-on:
-  pull_request:
-    paths:
-      - ".github/workflows/Integration.yml"
-      - "ext/**"
-      - "lib/**"
-      - "src/**"
-      - "test/**"
-      - "Project.toml"
-  push:
-    branches:
-      - main
-      - release-*
-    tags: ["*"]
-    paths:
-      - ".github/workflows/Integration.yml"
-      - "ext/**"
-      - "lib/**"
-      - "src/**"
-      - "test/**"
-      - "Project.toml"
-
-concurrency:
-  # Skip intermediate builds: always.
-  # Cancel intermediate builds: only if it is a pull request build.
-  group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
-
-jobs:
-  integration:
-    strategy:
-      fail-fast: false
-      matrix:
-        project:
-          - CUDA
-          - FancyArrayTypes
-          - FFT
-          - Flux
-          - Lux
-          - NNlib
-          - Python
-        version:
-          - "1.10"
-          - "1.11"
-        os:
-          - ubuntu-latest
-        # TODO: enable these tests
-        #   - linux-x86-ct6e-180-4tpu
-    uses: ./.github/workflows/CommonCI.yml
-    with:
-      julia_version: ${{ matrix.version }}
-      os: ${{ matrix.os }}
-      runtime: "both"
-      project: "test/integration/${{ matrix.project }}"
diff --git a/test/integration/Flux/Project.toml b/test/integration/Flux/Project.toml
index 28dd0fdfea..700f2c77d4 100644
--- a/test/integration/Flux/Project.toml
+++ b/test/integration/Flux/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

From b17f51605b162abc4437da0569910686b0bc5fe9 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 14:26:28 -0400
Subject: [PATCH 11/16] test: remove NNlib

---
 test/core/bcast.jl | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/test/core/bcast.jl b/test/core/bcast.jl
index 85f6f4249d..0b163fe526 100644
--- a/test/core/bcast.jl
+++ b/test/core/bcast.jl
@@ -1,8 +1,5 @@
-using Test
-using Reactant
-using Enzyme, NNlib
-using Reactant.MLIR
-using Statistics
+using Test, Reactant, Enzyme, Statistics
+using Reactant: MLIR
 
 @noinline function no(@nospecialize(x))
     x = @ccall $(Base.@cfunction(identity, Any, (Any,)))(x::Any)::Any

From 506a768ba9cb47e454860efac62519be82a08119 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 15:04:59 -0400
Subject: [PATCH 12/16] ci: set integration runtime

---
 .github/workflows/CommonCI.yml | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml
index 80484e8c3c..8ba66455d0 100644
--- a/.github/workflows/CommonCI.yml
+++ b/.github/workflows/CommonCI.yml
@@ -166,6 +166,15 @@ jobs:
           contents: |
             [Reactant]
             xla_runtime = "PJRT"
+      - name: "Setup Runtime Preferences for ${{ inputs.project }} (PJRT)"
+        if: ${{ (inputs.runtime == 'pjrt' || inputs.runtime == 'both') && inputs.project != '.' }}
+        uses: "DamianReeves/write-file-action@master"
+        with:
+          path: ${{ inputs.project }}/LocalPreferences.toml
+          write-mode: "overwrite"
+          contents: |
+            [Reactant]
+            xla_runtime = "PJRT"
 
       - name: "Run Tests (PJRT)"
         if: ${{ (inputs.runtime == 'pjrt' || inputs.runtime == 'both') && inputs.project == '.' }}
@@ -196,6 +205,15 @@ jobs:
           contents: |
             [Reactant]
             xla_runtime = "IFRT"
+      - name: "Setup Runtime Preferences for ${{ inputs.project }} (IFRT)"
+        if: ${{ (inputs.runtime == 'ifrt' || inputs.runtime == 'both') && inputs.project != '.' }}
+        uses: "DamianReeves/write-file-action@master"
+        with:
+          path: ${{ inputs.project }}/LocalPreferences.toml
+          write-mode: "overwrite"
+          contents: |
+            [Reactant]
+            xla_runtime = "IFRT"
 
       - name: "Run Tests (IFRT)"
         if: ${{ (inputs.runtime == 'ifrt' || inputs.runtime == 'both') && inputs.project == '.' }}

From 77982ab115568deb897f4e55812522046c7c259f Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 15:45:24 -0400
Subject: [PATCH 13/16] ci: better naming

---
 .github/workflows/CommonCI.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/CommonCI.yml b/.github/workflows/CommonCI.yml
index 8ba66455d0..57f6c1ab6d 100644
--- a/.github/workflows/CommonCI.yml
+++ b/.github/workflows/CommonCI.yml
@@ -190,7 +190,7 @@ jobs:
         env:
           ALLOW_RERESOLVE: ${{ !inputs.downgrade_testing }}
           REACTANT_TEST_GROUP: ${{ inputs.test_group }}
-      - name: "Run Integration Tests (PJRT)"
+      - name: "Run ${{ inputs.project }} Integration Tests (PJRT)"
         if: ${{ (inputs.runtime == 'pjrt' || inputs.runtime == 'both') && inputs.project != '.' }}
         timeout-minutes: 30
         run: julia --color=yes --project=${{ inputs.project }} --threads=auto --check-bounds=yes -O1 ${{ inputs.project }}/runtests.jl
@@ -229,7 +229,7 @@ jobs:
         env:
           ALLOW_RERESOLVE: ${{ !inputs.downgrade_testing }}
           REACTANT_TEST_GROUP: ${{ inputs.test_group }}
-      - name: "Run Integration Tests (IFRT)"
+      - name: "Run ${{ inputs.project }} Integration Tests (IFRT)"
         if: ${{ (inputs.runtime == 'ifrt' || inputs.runtime == 'both') && inputs.project != '.' }}
         timeout-minutes: 30
         run: julia --color=yes --project=${{ inputs.project }} --threads=auto --check-bounds=yes -O1 ${{ inputs.project }}/runtests.jl

From 81bcf1ff4a490e2c1ef94946c638f58bd245de38 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 15:48:58 -0400
Subject: [PATCH 14/16] test: fixes

---
 ext/ReactantFillArraysExt.jl       |  2 +-
 test/core/wrapped_arrays.jl        | 14 +-------------
 test/integration/NNlib/runtests.jl | 12 ++++++++++++
 3 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/ext/ReactantFillArraysExt.jl b/ext/ReactantFillArraysExt.jl
index f0c048581e..8ffc46bc77 100644
--- a/ext/ReactantFillArraysExt.jl
+++ b/ext/ReactantFillArraysExt.jl
@@ -1,6 +1,6 @@
 module ReactantFillArraysExt
 
-using Reactant: Reactant, TracedUtils, TracedRNumber, Ops, Sharding, unwrapped_eltype
+using Reactant: Reactant, TracedRNumber, Sharding, unwrapped_eltype
 using ReactantCore: ReactantCore
 using FillArrays: FillArrays, AbstractFill, Fill, Ones, Zeros, OneElement
 using GPUArraysCore: @allowscalar
diff --git a/test/core/wrapped_arrays.jl b/test/core/wrapped_arrays.jl
index a9b303eda5..d7d8dffe66 100644
--- a/test/core/wrapped_arrays.jl
+++ b/test/core/wrapped_arrays.jl
@@ -1,4 +1,4 @@
-using Reactant, Test, Statistics, NNlib, LinearAlgebra
+using Reactant, Test, Statistics, LinearAlgebra
 
 function view_getindex_1(x)
     x = view(x, 2:3, 1:2, :)
@@ -77,18 +77,6 @@ end
     @test v1 ≈ v2
 end
 
-function btranspose_badjoint(x)
-    x1 = NNlib.batched_transpose(x)
-    x2 = NNlib.batched_adjoint(x)
-    return x1 .+ x2
-end
-
-@testset "batched transpose/adjoint" begin
-    x = rand(4, 2, 3)
-    x_ra = Reactant.to_rarray(x)
-    @test @jit(btranspose_badjoint(x_ra)) ≈ btranspose_badjoint(x)
-end
-
 function bypass_permutedims(x)
     x = PermutedDimsArray(x, (2, 1, 3))  # Don't use permutedims here
     return view(x, 2:3, 1:2, :)
diff --git a/test/integration/NNlib/runtests.jl b/test/integration/NNlib/runtests.jl
index 876ee3eb9d..1436466886 100644
--- a/test/integration/NNlib/runtests.jl
+++ b/test/integration/NNlib/runtests.jl
@@ -792,3 +792,15 @@ end
     hlo = repr(@code_hlo(NNlib.gather(x_ra, idxs_ra)))
     @test !contains(hlo, "i64>")
 end
+
+function btranspose_badjoint(x)
+    x1 = NNlib.batched_transpose(x)
+    x2 = NNlib.batched_adjoint(x)
+    return x1 .+ x2
+end
+
+@testset "batched transpose/adjoint" begin
+    x = rand(4, 2, 3)
+    x_ra = Reactant.to_rarray(x)
+    @test @jit(btranspose_badjoint(x_ra)) ≈ btranspose_badjoint(x)
+end

From 86a46e8619f4f256f97694d688305332ad68a5a4 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Sun, 21 Sep 2025 17:20:26 -0400
Subject: [PATCH 15/16] ci: exclude extra tpu runs

---
 .github/workflows/CI.yml | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index bd86ebb2eb..e43ccaa786 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -117,6 +117,9 @@ jobs:
         os:
           - ubuntu-latest
           - linux-x86-ct6e-180-4tpu
+        exclude:
+          - os: linux-x86-ct6e-180-4tpu
+            version: "1.10"
         include:
           - project: CUDA
             os: linux-x86-a2-48-a100-4gpu

From 8065821bb46ae8cc5e7f7788d008cff3a53fd8e0 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Mon, 22 Sep 2025 10:02:19 -0400
Subject: [PATCH 16/16] test: use safetestset

---
 test/integration/CUDA/Project.toml | 1 +
 test/integration/CUDA/runtests.jl  | 6 +++---
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/test/integration/CUDA/Project.toml b/test/integration/CUDA/Project.toml
index f9363a2916..8b1c702ca1 100644
--- a/test/integration/CUDA/Project.toml
+++ b/test/integration/CUDA/Project.toml
@@ -2,6 +2,7 @@
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
diff --git a/test/integration/CUDA/runtests.jl b/test/integration/CUDA/runtests.jl
index ee37d9f6b8..9c4cbf88a3 100644
--- a/test/integration/CUDA/runtests.jl
+++ b/test/integration/CUDA/runtests.jl
@@ -1,11 +1,11 @@
-using Test
+using Test, SafeTestsets
 
 @testset "CUDA Integration" begin
-    @testset "CUDA" begin
+    @safetestset "CUDA" begin
         include("cuda.jl")
     end
 
-    @testset "KernelAbstractions" begin
+    @safetestset "KernelAbstractions" begin
         include("kernelabstractions.jl")
     end
 end