Parallel Testing + Distributed Docs build

.buildkite/pipeline.yml

@@ -173,7 +173,7 @@ steps:
             - src
             - ext
     command: |
-      julia --project --code-coverage=user --color=yes --threads=3 -e '
+      julia --project --code-coverage=user --color=yes --threads=auto -e '
         println("--- :julia: Instantiating project")
         using Pkg
         Pkg.instantiate()
@@ -185,6 +185,7 @@ steps:
         Pkg.activate("docs")
         include("docs/tutorials.jl")
         Pkg.activate("docs")
+        Pkg.instantiate()
         include("docs/make.jl")'
     agents:
       queue: "juliagpu"
@@ -197,6 +198,9 @@ steps:
     timeout_in_minutes: 240
 
 env:
+  RETESTITEMS_NWORKERS: 8
+  RETESTITEMS_NWORKER_THREADS: 2
+  LUX_DOCUMENTATION_NWORKERS: 4
   JULIA_AMDGPU_LOGGING_ENABLED: true
   SECRET_CODECOV_TOKEN: "jQ0BMTQgyZx7QGyU0Q2Ec7qB9mtE2q/tDu0FsfxvEG7/zOAGvXkyXrzIFFOQxvDoFcP+K2+hYZKMxicYdNqzr5wcxu505aNGN2GM3wyegAr+hO6q12bCFYx6qXzU9FLCCdeqINqn9gUSSOlGtWNFrbAlrTyz/D4Yo66TqBDzvaLL63FMnhCLaXW/zJt3hNuEAJaPY2O6Ze1rX2WZ3Y+i+s3uQ8aLImtoCJhPe8CRx+OhuYiTzGhynFfGntZ0738/1RN4gNM0S/hTC4gLE7XMVBanJpGh32rFaiDwW4zAyXKBrDkL3QA3MS1RvLTJxGJ085S16hCk0C4ddAhZCvIM9Q==;U2FsdGVkX1+bXdFeKMs5G79catOCyby2n07A2fg0FjVAvrjQLZ0yfvDS4paJiFikLkodho0khz2YALKb2Y0K6w=="
   SECRET_DOCUMENTER_KEY: "iRC4P/r5o9pARB670eK9jPlKQKgkTMDAyvp2GbLG8WwLuT8T1VcWx/o4+ofGlzbTh5Z+LuFgPXfgqkjGuoWLcocHNm78xQMNMywB4rcLB2shqp8xG2vhglgnTBBS4EiyPAtVqGyi5AKmfF95PfkJvnI0Lqg5P/RWQvNGywLAR0Ikgr/lqocm2CvkFGbpMzpGxGvj76JYOusVeKvGAp698TXqPabSZR2oZQLfYnEZnaO8ivkqvMGQSXfgzoIMjCOrN1rSa84SWeI9BDeBslzDHwaYGlvjpfCyviiLtKj4t5Acl1gVE0qxxZxWuALIU6z+C1W8TbW7ZDCBUFs6UTIT+Q==;U2FsdGVkX1+/HSgg1skLszz835vSO6mEtXMhG62ohQQUc5opdo7kEIAG2wCoJPQrqGyaF9kKDVvrN5G2MdjUyaLBYlv90RzXhjTiMNFdgI3M4K500xKq3itY/aEL7hUSMRKxTos8u4xhdbRboY4rPcqgtCJ2LHEjNxmml/NfEo/8lk291rGoEYQLTvKP9cuo4enmEVVRhqmabBzt1MDz0m4c8RufJWW2Ni4osaKRkYPjl/ijJ38wvRUZIiyCX7uofh+3iCKWn0111q5xFhn256Pm79Cx2ZP+yTp9sMsVNMJZ3UJ5r18F3H+zFHWWQSoiWpHn2WNB/2VUEyt0Lp1LnogKru96P2oYkXi6kqrA+qlLISUUU7R7ggJU0IRS6MjSGDyVzlaZG8m+RmY0bmQKrDwSeq1JMGkBpjwPY1o4yOnFRB7Rj1bzToLtd2IFSa8x0a2dUSyL5pBlyWklzZCxPp05R53RNSOi2KfhNfdZU2H7xEj5+z2aV5OidzowXIyYH8FlusMdk3NAOsvTbmBGiwvN4Zub9Exli06ZwARu/oJHLRh+hgOErIJ7DoX6nPrAtofSy6Etydpt+c4HkVZtGPWFSTMNWIGNx2NB1IfveOTU60H5emQ7zow5grXz4VTczqvCIh2hoQdSR4Oplr6+tDDLhtcGGHchHt473o2ygQ1m1tg7oSvMN7jmkUV1N6GniQofmlbr8d5LK4i/QtfC5GHCKIg3ohRlDvuvvKzvVWofgHX3NhXFTKK/CWAIp76iOaCWJcI562SpKyn+pFqYKpatJ42WfF3VbNpJYVMYMai5BwAE2RyZ6FhHbsaHq/NXO/dRJwHeDm4Pc/LFlGFdzpdbuf+w2DoePc56PlNmKsLNlZVlwbWcExKttI8nz3Th3aHNNtbIbD9awf1RdDspudQrTPWkyEopDVm7TkOj/J891U5p24PF5dasIJR19Tqpic3LVJuBXYRbL/Z79VRjeE3wBGLTDdhzJMA8TrS+yMSCF80bIw/F44o4WbA3Ya425mph9MIt/a137osRKATYqbustmVW/LfIyVhuHCOCRQsqTyFU+ff6Tp0EE2i1du90wosr+UutXiubYphCmuKkZONPbiXjpW1CAi40iAwxfgOVqAl13y4FlUp4EiGS7hPBUbvvEXMqT3ssfL+mlideH/v08PQCRcyG03zcCjCTmjXCggqHd+eEXhnsNZ4PFKCKiN+znR5SW+/p+kJTaBrX2e/kMU6kzjwb4NyNmZie0hHSneVtwJ1FuXJk/Zph4quv5KugCCx21xb5pePqxfKRW5jtW6r2Rc7OSNN4BHjwAcj8fOVV+12Ak7//o8mRh0aveYfoEvjCdaI8OPfjduDGfmzPUvXiqV9kGpovdlDUATyoVa3l1CowJ5r8KDOD6Ps89OG7TV2c7Wzxq2FQVjMFXxv/4wMZR1F/0zyH+ofPLVZjK3039z35GD4uoOW9Uc7WSr4FbxxuCDwOXWgstuk3rk6ASZFSe7RIwE/Y16d/aqzI+LG8pHqaEdhg6o6Y6JxBYNQo/JoglUOHwD+N5g5n9vfBNzf0xTlE/r0yjO3LCHyWzCnWr3QdKgzm6EDyL8GO+yQIbtXtw6lRQB/UEZ+ayt175r08Yhey95IsPwLVDFRRlG6pYwmzTlQOEwvqDI8SDMWboU+jp6a5jrbaAmqiIkaoiIzrV1QDp1x+Sqj0veqN+RtcpXLawJevz8dm76H+Mmp1br61nwvGcBaOKukICVj3iLeeu5tV5NoEJznWPwveHrcarZtKvOOeJbydmNAz286i0F1ocX337dt17jIkRv9sHbfqAVapob+eT7F3N/UY99GWGDVbXzaruQwsuPPR6MbLolG6buHQaKX3OZ/zJqGWfEAHw5yJKoKNe8aSgY2DsoITqPlbNRQQmOIMuF8ffD8L1stD/P5Ohth5Nql2W+l6y87/nqxkJ9y4FFS4QzrMrl9ztugfsRoYyeSWRydLUHlTCv155VsGAxjCMBQg1rP99Smfd02EbCFlWlypIw/zem0LZ1zVuz/Wjb03n+dzi2GIKRlTrt6YMrGGAcKI+3Pf1D0rsDhXNkdFUjOeofUkDbBr/splYCKLucDHFVdN88XyaQoj2fBymNJ4BqvK64TVOLwPGAQvh/rHZ5PkJR3lMI4fg+Kxdl9/5xDjkD9aV+yRvfqVGodNW/qofq34nrdb3co1tZ4BxtSANKdJg3Fv6U0I4DOMVsJTeOn/918M31rif0rKAwnHAkeyQVbZyEsFoqxvE8gUFs1zTRwZJWlmY0xnuVcM8pOh6hULeYGiF57ZlbvymygYqObe58YgrChRnF4NhKIIYzuz7mOSKRXqF3Cr0LNYHcktUH9wrqISxiHbaUQceYZ1D0q8UfiayeK9yppMkltcDUL9M93xjTGJK8pVzARXn6ETuEsNTtLvbU/KMDY7bnVc7n08suLCk1YeJB/sn0wuTbPt+27NeYIG1YXBEE0dsgJW4z64489h71v4xws856gFOHZx0L/nkW7l328HA3jltbgJFl52mQHAJwUZrt5sJef/k7gsTdX1zQtjKN8lFjo4qpvJUpenmO9nT+Wty5cjohlETBos8CdSqj4SjEu7/UhDt52evt33EayoWJ8TjKd4VRFYCXnM6eGnSMDqUU5f7DxVjrwHnT26jtq9ijKTiAxls7fYjN8TGT/S3CHZZAK1u5gSbWfkFOcE+mioboNwDvuvysjL6de+bsc7r35w4hLFnPmKemcde4pNQfEnuelBFJqwYZbcAkhN8AmtqIWPXBw9n3eUx/TJgMFEIoB/frNDRbB0WJKdBkjdE1NVvAUl3jDnZbWjG6rqE+6UvyGqKBpd0FRYAfg3ss3hVB70uluULKUBVazlNIQlqX+qYEMBXaDIkxcftre8KYebQyJnxiOB5V+eELvm6L28bK4Xh2tpXzJL7aDlQnL8dRNvQdZgDL62EXYhrc3mz0I/p7br3KMcnei/LaPRAgcsW7WKLwzE5id6JnpOJj4VXdkX7IUB4xQjDRsGKxhjbklMVFA8g/801khNlwzU/IoXsHBgTs7yZoFX/oo4Jyp514hwqPlvJEgci0OHiSA6Mx3le2nUh0SQH+AzFJ2vi7Bn1a4psiuqd+vJJ1iuNw5CBCZlV+GO8sG93BBGnLzZDoRvkIMbzwESFP3JYZ/lKs29CB2Adobl9YbwP3he0I9cD0A/RPC70gzTdVEfL6T4iPUhBr1Bn3YlUPeC2QvCTbpKkxDsfzchuq/y0xlmL4E7Rdb+4TSMlViXfnc6aoD9vvPMWLJFF2qrxRLKhUTse5V6RoE+EVmHSiX0Vd7sd/bYp7asOC0b1xL+zjfJ5DSrtMA/P8L1p+CoLNXgVfgzCB3sCa+GLSLS2INsL1Qtnfkl8IGaMDeV+VAyHjY0HCj0l1X99f/RzD6TYrZAkLS8h1EM/JjomglhVG9/HTKS20BBJeos5ifrVd38rhONJy0HCP28pn4rCIyIE4bNG+1tEsHAg4FDYgh/OYuBsaGYgha9TGV5lGIxmVCECq3IPpkPN1CsLqv3KuDvNeH6XOOAzVtFj4VoIV6QgRLP8+94ZiiEDaPQxQ7BZoqrqFYrxWHDtEuon46VtQ3Nfq/1Rq/HvszJv6JE77w7qvKlxG9sXgxzCDRqNrG83cwY2hpDBr8U0hPMrEx977Weja1aG/rG6uirNBcY5qAAOLDo+9RvV1xqvWFF8SkT97tzNUHbzw8tuUlCT9m4rshCG+jBw59rpUZwW+eR1ih9qU7Nyr3oNgi/zmkORF1duym8VSfW5dxtRBIqxxM0oSWoHti+HSd0VLdHw8jRpbQddMBr1sjD1jIgp3w2dU4oEthzStKCPY2/lAWBm+1Es1okGhEM3I939DRcYOjfJnTCtJLJ9DTKycVDMerXvHnCgImZ0Oh4mtLF+63hn+9wUc56owFeNqs+NJHqmBBFX2uNr/Rj9mzYkRRPsYYSyCB7jIS+Z8Zall6W3dwLcsE3uw/oPKx5bJDAhnp7kZgzLC0zlS2D0ZcNZuW2uUtwhZJM6OOyV+FUFgizmpIQAQ8Nm6n/1yk0asB4jZFf221a9ZmzvUfWKmmIR7OxX3qBH9x2uMMhemv9LZdEHMcjTeIXRYciMLWUNeWagYhDgV1cRBGCDTh2EhHvYX7ZXfpsHjLOR+sAEr7uR3siitf/mRkiLfT2YBgTACKKoj05UuC8aknEV4T5bWiye+gKGioml5G/fWYHyHow37g6D84n0cBTWmI0oPlg+rqpeRLOeYaTeCXOtM/7M1FHuGvzmBnag2vhKY2tpjVrg2nI3p4SRlzTyoQkyMfRXN87v5nAheVcLgrYtkv9aX7R6VMZ1UIsxn62ZHFa2IR6skB/xw7RRuJY5r5FIWs1LqIQDaon5L4C4v9rnBxMYoUM"

.github/workflows/CI.yml

@@ -37,7 +37,9 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         env:
-            GROUP: "CPU"
+          GROUP: "CPU"
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
         with:
           directories: src,ext

.github/workflows/Downgrade.yml

@@ -0,0 +1,41 @@
+name: Downgrade
+on:
+  pull_request:
+    branches:
+      - main
+    paths-ignore:
+      - 'docs/**'
+  push:
+    branches:
+      - master
+    paths-ignore:
+      - 'docs/**'
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        version: ['1.9']
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+      - uses: cjdoris/julia-downgrade-compat-action@v1
+        with:
+          skip: Pkg,TOML
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          GROUP: "CPU"
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
+      - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src,ext
+      - uses: codecov/codecov-action@v4
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          verbose: true
+          fail_ci_if_error: true

.github/workflows/Downstream.yml

@@ -58,6 +58,9 @@ jobs:
             @info "Not compatible with this release. No problem." exception=err
             exit(0)  # Exit immediately, as a success
           end
+        env:
+          RETESTITEMS_NWORKERS: 4
+          RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         with:

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.5.16"
+version = "0.5.17"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -49,35 +49,70 @@ LuxTrackerExt = "Tracker"
 LuxZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "0.1, 0.2"
-Adapt = "3, 4"
-ArrayInterface = "7"
-ChainRules = "1"
-ChainRulesCore = "1"
-ComponentArrays = "0.15.2"
+ADTypes = "0.2"
+Adapt = "4"
+Aqua = "0.8"
+ArrayInterface = "7.7"
+ChainRules = "1.62"
+ChainRulesCore = "1.21"
+ComponentArrays = "0.15.8"
 ConcreteStructs = "0.2"
 ConstructionBase = "1.5"
-FillArrays = "0.13, 1"
-Flux = "0.13, 0.14"
-Functors = "0.2, 0.3, 0.4"
-GPUArraysCore = "0.1"
-LinearAlgebra = "1"
-LuxAMDGPU = "0.1, 0.2"
-LuxCore = "0.1.6"
-LuxDeviceUtils = "0.1"
-LuxLib = "0.3"
-MacroTools = "0.5"
-Markdown = "1"
-Optimisers = "0.2, 0.3"
-PrecompileTools = "1"
-Random = "1"
+FillArrays = "1"
+Flux = "0.14.11"
+Functors = "0.4.4"
+GPUArraysCore = "0.1.6"
+LinearAlgebra = "1.9"
+Logging = "1.9"
+LuxAMDGPU = "0.2.2"
+LuxCUDA = "0.3.2"
+LuxCore = "0.1.7"
+LuxDeviceUtils = "0.1.14"
+LuxLib = "0.3.10"
+LuxTestUtils = "0.1.15"
+MacroTools = "0.5.13"
+Markdown = "1.9"
+Optimisers = "0.3"
+Pkg = "1.9"
+PrecompileTools = "1.2"
+Random = "1.9"
+ReTestItems = "1.23.1"
 Reexport = "1"
-ReverseDiff = "1"
-Setfield = "0.8, 1"
-SparseArrays = "1"
-Statistics = "1"
-Tracker = "0.2"
+ReverseDiff = "1.15"
+Setfield = "1"
+SparseArrays = "1.9"
+StableRNGs = "1"
+Statistics = "1.9"
+Test = "1.9"
+Tracker = "0.2.31"
 TruncatedStacktraces = "1.1"
 WeightInitializers = "0.1.5"
-Zygote = "0.6"
+Zygote = "0.6.69"
 julia = "1.9"
+
+[extras]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
+Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+LuxAMDGPU = "83120cb1-ca15-4f04-bf3b-6967d2e6b60b"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
+LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
+LuxTestUtils = "ac9de150-d08f-4546-94fb-7472b5760531"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[targets]
+test = ["ADTypes", "Aqua", "ChainRulesCore", "ComponentArrays", "Flux", "Functors", "LuxAMDGPU", "LuxCUDA", "LuxCore", "LuxLib", "LuxTestUtils", "Optimisers", "Random", "Setfield", "StableRNGs", "Statistics", "Test", "Zygote", "ReTestItems", "Reexport", "Logging"]

docs/Project.toml

@@ -2,6 +2,7 @@
 ANSIColoredPrinters = "a4c015fc-c6ff-483c-b24f-f7ea428134e9"
 Boltz = "4544d5e4-abc5-4dea-817f-29e4c205d9c8"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"

docs/src/introduction/index.md

@@ -2,7 +2,7 @@
 
 ## Installation
 
-Install [Julia v1.6 or above](https://julialang.org/downloads/). Lux.jl is available through
+Install [Julia v1.9 or above](https://julialang.org/downloads/). Lux.jl is available through
 the Julia package manager. You can enter it by pressing `]` in the REPL and then typing
 
 ```julia

docs/tutorials.jl

@@ -1,29 +1,38 @@
-using Literate
+using Distributed
 
-preprocess(path, str) = replace(str, "__DIR = @__DIR__" => "__DIR = \"$(dirname(path))\"")
+addprocs(parse(Int, get(ENV, "LUX_DOCUMENTATION_NWORKERS", "1")))
 
-get_example_path(p) = joinpath(@__DIR__, "..", "examples", p)
-OUTPUT = joinpath(@__DIR__, "src", "tutorials")
+@everywhere const LUX_DOCUMENTATION_NWORKERS = parse(Int,
+    get(ENV, "LUX_DOCUMENTATION_NWORKERS", "1"))
+@info "Lux Tutorial Build Running tutorials with $(LUX_DOCUMENTATION_NWORKERS) workers."
+@everywhere const CUDA_MEMORY_LIMIT = 100 ÷ LUX_DOCUMENTATION_NWORKERS
+
+@everywhere using Literate
+
+@everywhere function preprocess(path, str)
+    return replace(str, "__DIR = @__DIR__" => "__DIR = \"$(dirname(path))\"")
+end
+
+@everywhere get_example_path(p) = joinpath(@__DIR__, "..", "examples", p)
 
 BEGINNER_TUTORIALS = ["Basics/main.jl", "PolynomialFitting/main.jl", "SimpleRNN/main.jl"]
 INTERMEDIATE_TUTORIALS = ["NeuralODE/main.jl", "BayesianNN/main.jl", "HyperNet/main.jl"]
 ADVANCED_TUTORIALS = ["GravitationalWaveForm/main.jl"]
 
-withenv("JULIA_DEBUG" => "Literate") do
-    # tasks = []
-    for (d, paths) in (("beginner", BEGINNER_TUTORIALS),
-            ("intermediate", INTERMEDIATE_TUTORIALS),
-            ("advanced", ADVANCED_TUTORIALS)), (i, p) in enumerate(paths)
+TUTORIALS = [collect(Iterators.product(["beginner"], BEGINNER_TUTORIALS))...,
+    collect(Iterators.product(["intermediate"], INTERMEDIATE_TUTORIALS))...,
+    collect(Iterators.product(["advanced"], ADVANCED_TUTORIALS))...]
+
+pmap(enumerate(TUTORIALS)) do (i, (d, p))
+    println("Running tutorial $(i): $(p) on worker $(myid())")
+    withenv("JULIA_DEBUG" => "Literate",
+        "JULIA_CUDA_HARD_MEMORY_LIMIT" => "$(CUDA_MEMORY_LIMIT)%") do
         name = "$(i)_$(first(rsplit(p, "/")))"
         p_ = get_example_path(p)
-        # Literate.markdown(p_, joinpath(OUTPUT, d); execute=true, name=name, documenter=true,
-        #     preprocess=Base.Fix1(preprocess, p_))
-        # Using `@spawn` causes non-deterministic failures
-        jl_expr = "using Literate; preprocess(path, str) = replace(str, \"__DIR = @__DIR__\" => \"__DIR = \\\"\$(dirname(path))\\\"\"); Literate.markdown(\"$(p_)\", \"$(joinpath(OUTPUT, d))\"; execute=true, name=\"$name\", documenter=true, preprocess=Base.Fix1(preprocess, \"$(p_)\"))"
-        cm = `julia --project=$(@__DIR__) -e $(jl_expr)`
-        run(cm)
-        # task = Threads.@spawn run(cm)
-        # push!(tasks, task)
+        OUTPUT = joinpath(@__DIR__, "src", "tutorials")
+        res = Literate.markdown(p_, joinpath(OUTPUT, d); execute=true, name,
+            documenter=true, preprocess=Base.Fix1(preprocess, p_))
+        GC.gc(true)
+        return @isdefined(CUDA) && CUDA.reclaim()
     end
-    # return wait.(tasks)
 end

src/Lux.jl

@@ -13,8 +13,7 @@ PrecompileTools.@recompile_invalidations begin
 
     import LuxCore: AbstractExplicitLayer, AbstractExplicitContainerLayer,
                     initialparameters, initialstates, parameterlength, statelength,
-                    update_state,
-                    trainmode, testmode, setup, apply, display_name
+                    update_state, trainmode, testmode, setup, apply, display_name
     import LuxDeviceUtils: AbstractLuxDevice, AbstractLuxGPUDevice, AbstractLuxDeviceAdaptor
 end
 

src/chainrules.jl

@@ -21,9 +21,7 @@ function CRC.rrule(::typeof(merge), nt1::NamedTuple{F1}, nt2::NamedTuple{F2}) wh
         dnt2 = NamedTuple((f2 => getproperty(dy, f2) for f2 in F2))
         return (NoTangent(), dnt1, dnt2)
     end
-    function ∇merge(dy::Union{NoTangent, ZeroTangent})
-        return (NoTangent(), NoTangent(), NoTangent())
-    end
+    ∇merge(::Union{NoTangent, ZeroTangent}) = (NoTangent(), NoTangent(), NoTangent())
     return y, ∇merge
 end
 

src/deprecated.jl

@@ -11,9 +11,8 @@ This function has been deprecated. Use [`cpu_device`](@ref) instead.
 :::
 """
 function cpu(x)
-    Base.depwarn(
-        "`cpu` has been deprecated and will be removed in v0.6. Use `cpu_device` instead.",
-        :cpu)
+    Base.depwarn("`cpu` has been deprecated and will be removed in v0.6. Use \
+        `cpu_device` instead.", :cpu)
     return (cpu_device())(x)
 end
 
@@ -30,10 +29,11 @@ inside performance critical code will cause massive slowdowns due to type infere
 :::
 """
 function gpu(x)
-    @warn "Using `gpu` inside performance critical code will cause massive slowdowns due to type inference failure. Please update your code to use `gpu_device` API." maxlog=1
+    @warn "Using `gpu` inside performance critical code will cause massive slowdowns due \
+        to type inference failure. Please update your code to use `gpu_device` \
+        API." maxlog=1
 
-    Base.depwarn(
-        "`gpu` has been deprecated and will be removed in v0.6. Use `gpu_device` instead.",
-        :gpu)
+    Base.depwarn("`gpu` has been deprecated and will be removed in v0.6. Use \
+        `gpu_device` instead.", :gpu)
     return (gpu_device())(x)
 end

src/stacktraces.jl

@@ -13,46 +13,6 @@ function disable_stacktrace_truncation!(; disable::Bool=true)
     end
 end
 
-# NamedTuple / Tuples -- Lux uses them quite frequenty (states) making the error messages
-# too verbose
-@static if VERSION ≤ v"1.9" && !TruncatedStacktraces.DISABLE
-    function Base.show(io::IO, t::Type{<:Tuple})
-        if (TruncatedStacktraces.VERBOSE[] ||
-            !hasfield(t, :parameters) ||
-            length(t.parameters) == 0)
-            invoke(show, Tuple{IO, Type}, io, t)
-        else
-            try
-                fields = t.parameters
-                fields_truncated = length(fields) > 2 ? "$(fields[1]),$(fields[2]),…" :
-                                   (length(fields) == 2 ? "$(fields[1]),$(fields[2])" :
-                                    (length(fields) == 1 ? "$(fields[1])" : ""))
-                print(io, "Tuple{$fields_truncated}")
-            catch
-                invoke(show, Tuple{IO, Type}, io, t)
-            end
-        end
-    end
-
-    function Base.show(io::IO, t::Type{<:NamedTuple})
-        if (TruncatedStacktraces.VERBOSE[] ||
-            !hasfield(t, :parameters) ||
-            length(t.parameters) == 0)
-            invoke(show, Tuple{IO, Type}, io, t)
-        else
-            try
-                fields = first(t.parameters)
-                fields_truncated = length(fields) > 2 ? "$(fields[1]),$(fields[2]),…" :
-                                   (length(fields) == 2 ? "$(fields[1]),$(fields[2])" :
-                                    (length(fields) == 1 ? "$(fields[1])" : ""))
-                print(io, "NamedTuple{($fields_truncated),…}")
-            catch
-                invoke(show, Tuple{IO, Type}, io, t)
-            end
-        end
-    end
-end
-
 # Lux Layers
 ## layers/basic.jl
 @truncate_stacktrace WrappedFunction