ci: combine workflows

.buildkite/testing.yml

@@ -46,7 +46,7 @@ steps:
         env:
           BACKEND_GROUP: "CUDA"
         if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip ci\]/
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
         matrix:
           setup:
             julia:
@@ -101,7 +101,7 @@ steps:
           rocm: "*"
           rocmgpu: "*"
         if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip ci\]/
-        timeout_in_minutes: 60
+        timeout_in_minutes: 120
         matrix:
           setup:
             julia:

.github/workflows/CI.yml

@@ -33,22 +33,23 @@ jobs:
       matrix:
         version:
           - "1.10"
+          - "1"
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
         test_group:
           - "core_layers"
-          - "contrib"
-          - "helpers"
-          - "distributed"
           - "normalize_layers"
-          - "others"
-          - "autodiff"
           - "recurrent_layers"
-          - "eltype_match"
-          - "fluxcompat"
+          - "autodiff"
+          - "misc"
           - "reactant"
+        include:
+          - version: "1.10"
+            os: "macos-latest"
+            test_group: "all"
+          - version: "1.10"
+            os: "windows-latest"
+            test_group: "all"
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

.github/workflows/CIPreRelease.yml

@@ -37,15 +37,10 @@ jobs:
           - ubuntu-latest
         test_group:
           - "core_layers"
-          - "contrib"
-          - "helpers"
-          - "distributed"
           - "normalize_layers"
-          - "others"
-          - "autodiff"
           - "recurrent_layers"
-          - "eltype_match"
-          - "fluxcompat"
+          - "autodiff"
+          - "misc"
           - "reactant"
     steps:
       - uses: actions/checkout@v4
@@ -62,8 +57,30 @@ jobs:
             ${{ runner.os }}-test-${{ env.cache-name }}-
             ${{ runner.os }}-test-
             ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@v1
-      - uses: julia-actions/julia-runtest@v1
+      - name: "Install Dependencies"
+        run: |
+          import Pkg
+          dev_pkgs = Pkg.PackageSpec[]
+          for pkg in ("lib/LuxCore", "lib/MLDataDevices", "lib/WeightInitializers", "lib/LuxLib",)
+            push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
+          end
+          Pkg.develop(dev_pkgs)
+          Pkg.Registry.update()
+          Pkg.instantiate()
+          Pkg.activate("test")
+          dev_pkgs = Pkg.PackageSpec[]
+          for pkg in ("lib/LuxTestUtils", "lib/LuxLib", "lib/MLDataDevices", "lib/LuxCore", ".")
+            push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
+          end
+          Pkg.develop(dev_pkgs)
+          Pkg.instantiate()
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0}
+      - name: "Run Tests"
+        run: |
+          import Pkg, Lux
+          dir = dirname(pathof(Lux))
+          include(joinpath(dir, "../test/runtests.jl"))
+        shell: julia --color=yes --code-coverage=user --depwarn=yes --project=test {0}
         env:
           LUX_TEST_GROUP: ${{ matrix.test_group }}
           BACKEND_GROUP: "CPU"

.github/workflows/CI_LuxCore.yml

@@ -29,8 +29,6 @@ jobs:
           - "1"
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -53,7 +51,7 @@ jobs:
           Pkg.instantiate()
           Pkg.activate("lib/LuxCore/test")
           dev_pkgs = Pkg.PackageSpec[]
-          for pkg in ("lib/LuxCore",)
+          for pkg in ("lib/LuxCore", "lib/MLDataDevices")
             push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
           end
           Pkg.develop(dev_pkgs)
@@ -94,7 +92,7 @@ jobs:
           Pkg.instantiate()
           Pkg.activate("lib/LuxCore/test")
           dev_pkgs = Pkg.PackageSpec[]
-          for pkg in ("lib/LuxCore",)
+          for pkg in ("lib/LuxCore", "lib/MLDataDevices")
             push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
           end
           Pkg.develop(dev_pkgs)

.github/workflows/CI_LuxLib.yml

@@ -28,62 +28,53 @@ jobs:
       matrix:
         version:
           - "1.10"
+          - "1"
         os:
           - ubuntu-latest
         test_group:
           - "conv"
           - "dense"
-          - "batch_norm"
-          - "group_norm"
-          - "instance_norm"
-          - "layer_norm"
-          - "other_ops"
-          - "batched_ops"
-          - "others"
+          - "normalization"
+          - "misc"
         blas_backend:
           - "default"
         loopvec:
           - "true"
         include:
-          - os: ubuntu-latest
+          - version: "1.10"
+            os: ubuntu-latest
             test_group: "dense"
             blas_backend: "blis"
-            version: "1.10"
             loopvec: "true"
-          - os: ubuntu-latest
+          - version: "1.10"
+            os: ubuntu-latest
             test_group: "dense"
             blas_backend: "mkl"
-            version: "1.10"
             loopvec: "true"
-          - os: ubuntu-latest
+          - version: "1.10"
+            os: macos-latest
+            test_group: "dense"
+            blas_backend: "appleaccelerate"
+            loopvec: "true"
+          - version: "1.10"
+            os: ubuntu-latest
             test_group: "dense"
             blas_backend: "default"
-            version: "1.10"
-            loopvec: "false"
-          - os: ubuntu-latest
-            test_group: "batched_ops"
-            blas_backend: "default"
-            version: "1.10"
             loopvec: "false"
-          - os: ubuntu-latest
-            test_group: "other_ops"
+          - version: "1.10"
+            os: ubuntu-latest
+            test_group: "misc"
             blas_backend: "default"
-            version: "1.10"
             loopvec: "false"
-          - os: macos-latest
-            test_group: "dense"
-            blas_backend: "appleaccelerate"
-            version: "1.10"
-            loopvec: "true"
-          - os: macos-latest
+          - version: "1.10"
+            os: macos-latest
             test_group: "all"
             blas_backend: "default"
-            version: "1.10"
             loopvec: "true"
-          - os: windows-latest
+          - version: "1.10"
+            os: windows-latest
             test_group: "all"
             blas_backend: "default"
-            version: "1.10"
             loopvec: "true"
     steps:
       - uses: actions/checkout@v4
@@ -112,7 +103,7 @@ jobs:
           Pkg.instantiate()
           Pkg.activate("lib/LuxLib/test")
           dev_pkgs = Pkg.PackageSpec[]
-          for pkg in ("lib/LuxTestUtils", "lib/LuxLib")
+          for pkg in ("lib/LuxTestUtils", "lib/LuxLib", "lib/MLDataDevices")
             push!(dev_pkgs, Pkg.PackageSpec(path=pkg))
           end
           Pkg.develop(dev_pkgs)
@@ -142,21 +133,6 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
-      matrix:
-        test_group:
-          - "conv"
-          - "dense"
-          - "batch_norm"
-          - "group_norm"
-          - "instance_norm"
-          - "layer_norm"
-          - "other_ops"
-          - "batched_ops"
-          - "others"
-        blas_backend:
-          - "default"
-        loopvec:
-          - "true"
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -189,9 +165,9 @@ jobs:
           include(joinpath(dir, "../test/runtests.jl"))
         shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/LuxLib/test {0}
         env:
-          LUXLIB_TEST_GROUP: ${{ matrix.test_group }}
-          LUXLIB_BLAS_BACKEND: ${{ matrix.blas_backend }}
-          LUXLIB_LOAD_LOOPVEC: ${{ matrix.loopvec }}
+          LUXLIB_TEST_GROUP: "all"
+          LUXLIB_BLAS_BACKEND: "default"
+          LUXLIB_LOAD_LOOPVEC: "true"
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: lib/LuxLib/src,lib/LuxLib/ext,lib/LuxCore/src,lib/LuxCore/ext,lib/MLDataDevices/src,lib/MLDataDevices/ext,lib/LuxTestUtils/src

.github/workflows/CI_LuxTestUtils.yml

@@ -28,8 +28,6 @@ jobs:
           - "1"
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

.github/workflows/CI_WeightInitializers.yml

@@ -28,8 +28,6 @@ jobs:
           - "1"
         os:
           - ubuntu-latest
-          - macos-latest
-          - windows-latest
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

Project.toml

@@ -76,10 +76,10 @@ Compat = "4.15"
 ComponentArrays = "0.15.16"
 ConcreteStructs = "0.2.3"
 DispatchDoctor = "0.4.12"
-Enzyme = "0.13.1"
-EnzymeCore = "0.8.1"
+Enzyme = "0.13.13"
+EnzymeCore = "0.8.5"
 FastClosures = "0.3.2"
-Flux = "0.14.20"
+Flux = "0.14.25"
 ForwardDiff = "0.10.36"
 FunctionWrappers = "1.1.3"
 Functors = "0.4.12"

docs/Project.toml

@@ -36,7 +36,7 @@ ChainRulesCore = "1.24"
 ComponentArrays = "0.15"
 Documenter = "1.4"
 DocumenterVitepress = "0.1.3"
-Enzyme = "0.13"
+Enzyme = "0.13.13"
 FiniteDiff = "2.23.1"
 ForwardDiff = "0.10.36"
 Functors = "0.4.12"
@@ -48,7 +48,7 @@ Lux = "1"
 LuxCUDA = "0.3.2"
 LuxCore = "1"
 LuxLib = "1.3.4"
-LuxTestUtils = "1.2"
+LuxTestUtils = "1.4"
 MLDataDevices = "1.4"
 Optimisers = "0.3.3"
 Pkg = "1.10"

docs/src/manual/compiling_lux_models.md

@@ -21,6 +21,11 @@ using Lux, Reactant, Enzyme, Random, Zygote
 using Functors, Optimisers, Printf
 ```
 
+!!! tip "Running on alternate accelerators"
+
+    `Reactant.set_default_backend("gpu")` sets the default backend to CUDA and
+    `Reactant.set_default_backend("tpu")` sets the default backend to TPU.
+
 !!! tip "Using the `TrainState` API"
 
     If you are using the [`Training.TrainState`](@ref) API, skip to the
@@ -149,15 +154,12 @@ function train_model(model, ps, st, dataloader)
     train_state = Training.TrainState(model, ps, st, Adam(0.001f0))
 
     for iteration in 1:1000
-        for (xᵢ, yᵢ) in dataloader
-            grads, loss, stats, train_state = Training.single_train_step!(
+        for (i, (xᵢ, yᵢ)) in enumerate(dataloader)
+            _, loss, _, train_state = Training.single_train_step!(
                 AutoEnzyme(), MSELoss(), (xᵢ, yᵢ), train_state)
-        end
-        if iteration % 100 == 0 || iteration == 1
-            # We need to do this since scalar outputs are currently expressed as a zero-dim
-            # array
-            loss = Array(loss)[]
-            @printf("Iter: [%4d/%4d]\tLoss: %.8f\n", iteration, 1000, loss)
+            if (iteration % 100 == 0 || iteration == 1) && i == 1
+                @printf("Iter: [%4d/%4d]\tLoss: %.8f\n", iteration, 1000, loss)
+            end
         end
     end
 

lib/LuxCore/Project.toml

@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.1.0"
+version = "1.1.1"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -34,11 +34,11 @@ ArrayInterface = "7.9"
 ChainRulesCore = "1.24"
 Compat = "4.15.0"
 DispatchDoctor = "0.4.10"
-EnzymeCore = "0.7.7, 0.8"
+EnzymeCore = "0.8.5"
 Functors = "0.4.12"
 MLDataDevices = "1"
 Random = "1.10"
-Reactant = "0.2.3"
+Reactant = "0.2.4"
 ReverseDiff = "1.15"
 Setfield = "1"
 Tracker = "0.2.34"

lib/LuxCore/test/Project.toml

@@ -11,7 +11,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Aqua = "0.8.7"
-EnzymeCore = "0.7.7"
+EnzymeCore = "0.8.5"
 ExplicitImports = "1.9.0"
 Functors = "0.4.12"
 MLDataDevices = "1.0.0"

lib/LuxLib/Project.toml

@@ -65,8 +65,8 @@ ChainRulesCore = "1.24"
 Compat = "4.15.0"
 CpuId = "0.3"
 DispatchDoctor = "0.4.12"
-Enzyme = "0.13.1"
-EnzymeCore = "0.8.1"
+Enzyme = "0.13.13"
+EnzymeCore = "0.8.5"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.36"
 Hwloc = "3.2"