Merge branch 'main' into ap/up_opt

.buildkite/testing.yml

@@ -25,9 +25,10 @@ steps:
             julia:
               - "1"
             group:
-              - "!fluxcompat,distributed"
+              - "!fluxcompat,distributed,recurrent_layers"
               - "fluxcompat"
               - "distributed"
+              - "recurrent_layers"
 
   - group: ":telescope: Downstream CUDA"
     steps:
@@ -81,9 +82,10 @@ steps:
             julia:
               - "1"
             group:
-              - "!fluxcompat,distributed"
+              - "!fluxcompat,distributed,recurrent_layers"
               - "fluxcompat"
               - "distributed"
+              - "recurrent_layers"
 
   - group: ":telescope: Downstream AMD GPU"
     steps:

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.0.5"
+version = "1.0.6"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -64,7 +64,7 @@ LuxTrackerExt = "Tracker"
 LuxZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "1.5"
+ADTypes = "1.8.1"
 Adapt = "4"
 ArgCheck = "2.3"
 ArrayInterface = "7.9"
@@ -74,25 +74,25 @@ Compat = "4.15"
 ComponentArrays = "0.15.16"
 ConcreteStructs = "0.2.3"
 DispatchDoctor = "0.4.12"
-Enzyme = "0.12.26"
-EnzymeCore = "0.7.7"
+Enzyme = "0.13.1"
+EnzymeCore = "0.8.1"
 FastClosures = "0.3.2"
-Flux = "0.14.17"
+Flux = "0.14.20"
 ForwardDiff = "0.10.36"
 FunctionWrappers = "1.1.3"
 Functors = "0.4.12"
 GPUArraysCore = "0.1.6"
 LinearAlgebra = "1.10"
 LossFunctions = "0.11.1"
 LuxCore = "1"
-LuxLib = "1.2"
+LuxLib = "1.3"
 MLDataDevices = "1.1"
 MLUtils = "0.4.4"
 MPI = "0.20.19"
 MacroTools = "0.5.13"
 Markdown = "1.10"
 NCCL = "0.1.1"
-NNlib = "0.9.21"
+NNlib = "0.9.24"
 Optimisers = "0.3.3"
 Preferences = "1.4.3"
 Random = "1.10"

docs/run_single_tutorial.jl

@@ -8,15 +8,33 @@ name = ARGS[1]
 pkg_log_path = joinpath(storage_dir, "$(name)_pkg.log")
 output_directory = ARGS[2]
 path = ARGS[3]
+push!(LOAD_PATH, "@literate")  # Should have the Literate and InteractiveUtils packages
 
 io = open(pkg_log_path, "w")
-Pkg.develop(; path=joinpath(@__DIR__, ".."), io)
+warn_old_version = try
+    Pkg.develop(; path=joinpath(@__DIR__, ".."), io)
+    false
+catch err
+    err isa Pkg.Resolve.ResolverError || rethrow()
+    @warn "Failed to install the latest version of Lux.jl. This is possible when the \
+           downstream packages haven't been updated to support latest releases yet." err=err
+    true
+end
 Pkg.instantiate(; io)
 close(io)
 
 using Literate
 
 function preprocess(path, str)
+    if warn_old_version
+        str = """
+        # !!! danger "Using older version of Lux.jl"
+
+        #     This tutorial cannot be run on the latest Lux.jl release due to downstream
+        #     packages not being updated yet.
+
+        \n\n""" * str
+    end
     new_str = replace(str, "__DIR = @__DIR__" => "__DIR = \"$(dirname(path))\"")
     appendix_code = """
     # ## Appendix

docs/tutorials.jl

@@ -55,6 +55,8 @@ const NTASKS = min(
 
 @info "Starting Lux Tutorial Build with $(NTASKS) tasks."
 
+run(`$(Base.julia_cmd()) --startup=no --code-coverage=user --threads=$(Threads.nthreads()) --project=@literate -e 'import Pkg; Pkg.add(["Literate", "InteractiveUtils"])'`)
+
 asyncmap(TUTORIALS_BUILDING; ntasks=NTASKS) do (i, (d, p))
     @info "Running Tutorial $(i): $(p) on task $(current_task())"
     path = joinpath(@__DIR__, "..", "examples", p)
@@ -66,8 +68,7 @@ asyncmap(TUTORIALS_BUILDING; ntasks=NTASKS) do (i, (d, p))
     withenv("JULIA_NUM_THREADS" => "$(Threads.nthreads())",
         "JULIA_CUDA_HARD_MEMORY_LIMIT" => "$(100 ÷ NTASKS)%",
         "JULIA_PKG_PRECOMPILE_AUTO" => "0", "JULIA_DEBUG" => "Literate") do
-        cmd = `$(Base.julia_cmd()) --code-coverage=user --threads=$(Threads.nthreads()) --project=$(tutorial_proj) "$(file)" "$(name)" "$(output_directory)" "$(path)"`
-        run(cmd)
+        run(`$(Base.julia_cmd()) --startup=no --code-coverage=user --threads=$(Threads.nthreads()) --project=$(tutorial_proj) "$(file)" "$(name)" "$(output_directory)" "$(path)"`)
     end
 
     return

examples/Basics/Project.toml

@@ -1,8 +1,6 @@
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -14,7 +12,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 CUDA_Driver_jll = "0.9, 0.10"
 ComponentArrays = "0.15"
 ForwardDiff = "0.10"
-Literate = "2"
 Lux = "1"
 LuxCUDA = "0.3"
 Optimisers = "0.3"

examples/BayesianNN/Project.toml

@@ -1,9 +1,7 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
@@ -14,7 +12,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 CairoMakie = "0.12"
 Functors = "0.4"
 LinearAlgebra = "1"
-Literate = "2"
 Lux = "1"
 Random = "1"
 Tracker = "0.2"

examples/GravitationalWaveForm/Project.toml

@@ -1,9 +1,7 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
@@ -16,7 +14,6 @@ SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 CairoMakie = "0.12"
 ComponentArrays = "0.15"
 LineSearches = "7"
-Literate = "2"
 Lux = "1"
 Optimization = "4"
 OptimizationOptimJL = "0.4"

examples/HyperNet/Project.toml

@@ -1,8 +1,6 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
@@ -18,7 +16,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 ADTypes = "1"
 ComponentArrays = "0.15"
-Literate = "2"
 Lux = "1"
 LuxCUDA = "0.3"
 MLDatasets = "0.7"

examples/NeuralODE/Project.toml

@@ -1,7 +1,6 @@
 [deps]
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
@@ -17,7 +16,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ComponentArrays = "0.15"
-Literate = "2"
 Lux = "1"
 LuxCUDA = "0.2, 0.3"
 MLDatasets = "0.5, 0.7"

examples/OptimizationIntegration/Project.toml

@@ -1,8 +1,6 @@
 [deps]
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
@@ -17,8 +15,6 @@ SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 [compat]
 CairoMakie = "0.12.10"
 ComponentArrays = "0.15.17"
-InteractiveUtils = "<0.0.1, 1"
-Literate = "2.19"
 Lux = "1"
 LuxCUDA = "0.3.3"
 MLUtils = "0.4.4"

examples/PINN2DPDE/Project.toml

@@ -1,8 +1,6 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
@@ -14,10 +12,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ADTypes = "1.7.1"
+ADTypes = "1.8.1"
 CairoMakie = "0.12.10"
-InteractiveUtils = "<0.0.1, 1"
-Literate = "2"
 Lux = "1"
 LuxCUDA = "0.3.3"
 MLUtils = "0.4.4"

examples/PolynomialFitting/Project.toml

@@ -1,8 +1,6 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -14,7 +12,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [compat]
 ADTypes = "1"
 CairoMakie = "0.12"
-Literate = "2"
 Lux = "1"
 LuxCUDA = "0.3"
 Optimisers = "0.3"

examples/SimpleChains/Project.toml

@@ -1,7 +1,5 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
@@ -14,7 +12,6 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ADTypes = "1"
-Literate = "2"
 Lux = "1"
 MLDatasets = "0.7.14"
 MLUtils = "0.4"

examples/SimpleRNN/Project.toml

@@ -1,8 +1,6 @@
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
@@ -14,8 +12,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 ADTypes = "1"
-JLD2 = "0.4, 0.5"
-Literate = "2"
+JLD2 = "0.5"
 Lux = "1"
 LuxCUDA = "0.3"
 MLUtils = "0.4"

src/helpers/losses.jl

@@ -57,7 +57,7 @@ function CRC.rrule(::typeof(fused_agg), ::typeof(sum), lfn::LossFunctions.Traits
 end
 
 function EnzymeRules.augmented_primal(
-        cfg::EnzymeRules.ConfigWidth{1}, func::EnzymeCore.Const{typeof(fused_agg)},
+        cfg::EnzymeRules.RevConfigWidth{1}, func::EnzymeCore.Const{typeof(fused_agg)},
         ::Type{<:EnzymeCore.Active}, agg_f::EnzymeCore.Const{typeof(sum)},
         lfn::EnzymeCore.Const{<:LossFunctions.Traits.Loss},
         x::EnzymeCore.Annotation{<:AbstractArray}, y::EnzymeCore.Const)
@@ -71,7 +71,7 @@ function EnzymeRules.augmented_primal(
 end
 
 function EnzymeRules.reverse(
-        cfg::EnzymeRules.ConfigWidth{1}, ::EnzymeCore.Const{typeof(fused_agg)},
+        cfg::EnzymeRules.RevConfigWidth{1}, ::EnzymeCore.Const{typeof(fused_agg)},
         dret::EnzymeCore.Active, (cache_x, cache_y), agg_f::EnzymeCore.Const{typeof(sum)},
         lfn::EnzymeCore.Const{<:LossFunctions.Traits.Loss},
         x::EnzymeCore.Annotation{<:AbstractArray}, y::EnzymeCore.Const)

test/Project.toml

@@ -39,15 +39,15 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-ADTypes = "1.5"
+ADTypes = "1.8.1"
 Adapt = "4"
 Aqua = "0.8.4"
 CUDA_Driver_jll = "0.9, 0.10"
 ChainRulesCore = "1.24"
 ComponentArrays = "0.15.16"
 DispatchDoctor = "0.4.12"
 Documenter = "1.4"
-Enzyme = "0.12.26"
+Enzyme = "0.13.1"
 ExplicitImports = "1.9.0"
 ForwardDiff = "0.10.36"
 Functors = "0.4.12"
@@ -56,11 +56,11 @@ InteractiveUtils = "<0.0.1, 1"
 LinearAlgebra = "1.10"
 Logging = "1.10"
 LuxCore = "1.0"
-LuxLib = "1.0"
-LuxTestUtils = "1.2"
+LuxLib = "1.3"
+LuxTestUtils = "1.3"
 MLDataDevices = "1.1"
 MLUtils = "0.4.3"
-NNlib = "0.9.21"
+NNlib = "0.9.24"
 OneHotArrays = "0.2.5"
 Optimisers = "0.3.3"
 Pkg = "1.10"

test/contrib/freeze_tests.jl

@@ -35,7 +35,8 @@
             @jet m(x, ps_c, st)
 
             __f = (x, ps) -> sum(first(m(x, ps, st)))
-            @test_gradients(__f, x, ps_c; atol=1.0f-3, rtol=1.0f-3)
+            @test_gradients(__f, x, ps_c; atol=1.0f-3, rtol=1.0f-3,
+                enzyme_set_runtime_activity=true)
         end
 
         @testset "LuxDL/Lux.jl#427" begin
@@ -84,7 +85,8 @@ end
 
         @jet fd(x, ps, st)
         __f = (x, ps) -> sum(first(fd(x, ps, st)))
-        @test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3)
+        @test_gradients(__f, x, ps; atol=1.0f-3, rtol=1.0f-3,
+            enzyme_set_runtime_activity=true)
 
         fd = Lux.Experimental.freeze(d, ())
         @test fd === d

test/helpers/compact_tests.jl

@@ -316,7 +316,7 @@
                 @compact(W=randn(d_out, d_in), b=zeros(d_out), incr=1) do x
                     y = W * x
                     incr *= 10
-                    return act.(y .+ b) .+ incr
+                    @return act.(y .+ b) .+ incr
                 end
             end
 
@@ -329,12 +329,7 @@
             @test st_new.incr == 10
             _, st_new = model(x, ps, st_new)
             @test st_new.incr == 100
-
-            # By default creates a closure so type cannot be inferred
-            inf_type = Core.Compiler._return_type(
-                model, Tuple{typeof(x), typeof(ps), typeof(st)}).parameters
-            @test inf_type[1] === Any
-            @test inf_type[2] === NamedTuple
+            @test @inferred(model(x, ps, st)) isa Any
 
             function ScaledDense2(; d_in=5, d_out=7, act=relu)
                 @compact(W=randn(d_out, d_in), b=zeros(d_out), incr=1) do x