# Introduction {#getting-started} ## Installation {#Installation} Install [Julia v1.10 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 `add Lux`. Alternatively, you can also do ```julia import Pkg Pkg.add("Lux") ``` ::: tip Update to v1 If you are using a pre-v1 version of Lux.jl, please see the [Updating to v1 section](/introduction/updating_to_v1#updating-to-v1) for instructions on how to update. ::: ## Quickstart {#Quickstart} ::: tip Pre-Requisites You need to install `Optimisers`, `Reactant` and `Enzyme` if not done already. `Pkg.add(["Optimisers", "Enzyme", "Reactant"])` ::: ```julia using Lux, Random, Optimisers, Enzyme, Reactant ``` ```ansi ┌ Debug: Persistent compilation cache enabled. Using base directory: /home/runner/.julia/scratchspaces/3c362404-f566-11ee-1572-e11a4b42c853/xla_persistent_cache_0_0_383 └ @ Reactant.PersistentCompileCache ~/.julia/packages/Reactant/auube/src/PersistentCompileCache.jl:28 ┌ Debug: Kernel cache enabled: false └ @ Reactant.PersistentCompileCache ~/.julia/packages/Reactant/auube/src/PersistentCompileCache.jl:33 ┌ Debug: Autotune cache enabled: true └ @ Reactant.PersistentCompileCache ~/.julia/packages/Reactant/auube/src/PersistentCompileCache.jl:38 ┌ Debug: Registering Backend cpu with Priority 100. └ @ Reactant.Accelerators.Registration ~/.julia/packages/Reactant/auube/src/accelerators/Registration.jl:39 ┌ Debug: Using libReactantExtra from /home/runner/.julia/artifacts/51b5a70aa846cc96fcb405e27b836039a06597fd/lib/libReactantExtra.so └ @ Reactant.XLA ~/.julia/packages/Reactant/auube/src/xla/XLA.jl:200 ┌ Debug: REACTANT_XLA_RUNTIME: │  REACTANT_XLA_RUNTIME = "PJRT" └ @ Reactant.XLA ~/.julia/packages/Reactant/auube/src/xla/XLA.jl:237 ``` We take randomness very seriously ```julia # Seeding rng = Random.default_rng() Random.seed!(rng, 0) ``` ```ansi Random.TaskLocalRNG() ``` Build the model ```julia # Construct the layer model = Chain(Dense(128, 256, tanh), Chain(Dense(256, 256, tanh), Dense(256, 10))) ``` ```ansi Chain( layer_1 = Dense(128 => 256, tanh), # 33_024 parameters layer_2 = Chain( layer_1 = Dense(256 => 256, tanh), # 65_792 parameters layer_2 = Dense(256 => 10), # 2_570 parameters ), )  # Total: 101_386 parameters,  # plus 0 states. ``` Models don't hold parameters and states so initialize them. From there on, we can just use our standard AD and Optimisers API. However, here we will show how to use Lux's Training API that provides an uniform API over all supported AD systems. ```julia # Get the device determined by Lux dev = reactant_device() # Parameter and State Variables ps, st = Lux.setup(rng, model) |> dev # Dummy Input x = rand(rng, Float32, 128, 2) |> dev # Run the model ## We need to use @jit to compile and run the model with Reactant y, st = @jit Lux.apply(model, x, ps, st) ## For best performance, first compile the model with Reactant and then run it apply_compiled = @compile Lux.apply(model, x, ps, st) apply_compiled(model, x, ps, st) # Gradients ## First construct a TrainState train_state = Training.TrainState(model, ps, st, Adam(0.0001f0)) ## We can compute the gradients using Training.compute_gradients ## TrainState handles compilation internally gs, loss, stats, train_state = Lux.Training.compute_gradients( AutoEnzyme(), MSELoss(), (x, dev(rand(rng, Float32, 10, 2))), train_state ) ## Optimization train_state = Training.apply_gradients!(train_state, gs) # or Training.apply_gradients (no `!` at the end) # Both these steps can be combined into a single call (preferred approach) gs, loss, stats, train_state = Training.single_train_step!( AutoEnzyme(), MSELoss(), (x, dev(rand(rng, Float32, 10, 2))), train_state ) ``` ```ansi ((layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.014252576 -0.01007682 … -0.0066116317 -0.012908219; -0.0066900905 0.001967808 … -0.003597645 -0.0010663884; … ; -0.013411599 -0.009727624 … -0.0062034056 -0.012329484; -0.021803929 -0.008291057 … -0.010640304 -0.014436427]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-0.016153783, -0.0064898166, 0.0040475307, 0.048172455, 0.018123386, -0.026255043, 0.031807806, -0.00242692, 0.0047841636, 0.030119237 … -0.018572997, 0.002068058, 0.009790083, 0.018018095, -0.01424905, -0.060925536, -0.016411005, -0.014629741, -0.015240658, -0.023550106])), layer_2 = (layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-2.1789727f-5 -3.8227197f-5 … 2.2419797f-6 -6.324912f-5; 0.008564907 0.006512975 … -0.006969174 0.009865682; … ; -0.021595063 -0.021455316 … 0.013971797 -0.033741944; -0.023399103 -0.018417688 … 0.018593067 -0.028052688]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-8.567904f-5, 0.015442661, -0.012886962, 0.052984916, 0.000940611, -0.013805261, -0.025449416, 0.012329709, 0.09247952, -0.038467444 … 0.02124564, -0.012647957, -0.00045393303, -0.026561411, -0.025026277, -0.044307135, 0.033703856, 0.02700105, -0.049719024, -0.043526463])), layer_2 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.041988406 0.0035988672 … 0.012094588 -0.006939848; -0.06445931 0.02036321 … 0.044551518 0.008209803; … ; 0.24316937 -0.04692524 … -0.11571932 0.0070323865; 0.027664784 -0.008376195 … -0.018484466 -0.0030616056]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.04771936, 0.079492785, -0.063023046, 0.011547155, -0.4143036, -0.118838035, -0.16815118, -0.2632838, -0.2873201, -0.0339642])))), Reactant.ConcretePJRTNumber{Float32, 1}(1.0014447f0), NamedTuple(), Lux.Training.TrainState{Lux.Training.TrainingBackendCache{Lux.Training.ReactantBackend{Static.True, Missing, Nothing, AutoEnzyme{Nothing, Nothing}}, Static.False, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{compiled_gradient_function::Reactant.Compiler.Thunk{typeof(ReactantExt.compute_gradients_internal), Symbol("##compute_gradients_internal_reactant#433"), false, Tuple{GenericLossFunction{typeof(Lux.LossFunctionImpl.l2_distance_loss), typeof(Statistics.mean)}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}}, Nothing}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}}, update_function::Reactant.Compiler.Thunk{typeof(Optimisers.update!), Symbol("##update!_reactant#686"), false, Tuple{@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}}, compiled_grad_and_step_function::Reactant.Compiler.Thunk{typeof(ReactantExt.compute_gradients_internal_and_step!), Symbol("##compute_gradients_internal_and_step!_reactant#1007"), false, Tuple{GenericLossFunction{typeof(Lux.LossFunctionImpl.l2_distance_loss), typeof(Statistics.mean)}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}}, Nothing}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}, @NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, Bool}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}}, is_sharded::Bool}}, GenericLossFunction{typeof(Lux.LossFunctionImpl.l2_distance_loss), typeof(Statistics.mean)}, Nothing, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}}, Nothing}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}, Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, @NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}}}}(Lux.Training.TrainingBackendCache{Lux.Training.ReactantBackend{Static.True, Missing, Nothing, AutoEnzyme{Nothing, Nothing}}, Static.False, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{compiled_gradient_function::Reactant.Compiler.Thunk{typeof(ReactantExt.compute_gradients_internal), Symbol("##compute_gradients_internal_reactant#433"), false, Tuple{GenericLossFunction{typeof(Lux.LossFunctionImpl.l2_distance_loss), typeof(Statistics.mean)}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}}, Nothing}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}}, update_function::Reactant.Compiler.Thunk{typeof(Optimisers.update!), Symbol("##update!_reactant#686"), false, Tuple{@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}}, compiled_grad_and_step_function::Reactant.Compiler.Thunk{typeof(ReactantExt.compute_gradients_internal_and_step!), Symbol("##compute_gradients_internal_and_step!_reactant#1007"), false, Tuple{GenericLossFunction{typeof(Lux.LossFunctionImpl.l2_distance_loss), typeof(Statistics.mean)}, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}}, Nothing}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, @NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{layer_1::@NamedTuple{}, layer_2::@NamedTuple{}}}, @NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}, layer_2::@NamedTuple{weight::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 2, 1}, Reactant.ConcretePJRTArray{Float32, 2, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}, bias::Optimisers.Leaf{Lux.ReactantCompatibleOptimisers.ReactantOptimiser{Adam{Reactant.ConcretePJRTNumber{Float32, 1}, Tuple{Reactant.ConcretePJRTNumber{Float64, 1}, Reactant.ConcretePJRTNumber{Float64, 1}}, Reactant.ConcretePJRTNumber{Float64, 1}}}, Tuple{Reactant.ConcretePJRTArray{Float32, 1, 1}, Reactant.ConcretePJRTArray{Float32, 1, 1}, Tuple{Reactant.ConcretePJRTNumber{Float32, 1}, Reactant.ConcretePJRTNumber{Float32, 1}}}}}}}, @NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{layer_1::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}, layer_2::@NamedTuple{weight::Reactant.ConcretePJRTArray{Float32, 2, 1}, bias::Reactant.ConcretePJRTArray{Float32, 1, 1}}}}, Bool}, Reactant.XLA.PJRT.LoadedExecutable, Reactant.XLA.PJRT.Device, Reactant.XLA.PJRT.Client, Tuple{}, Vector{Bool}}, is_sharded::Bool}}(Lux.Training.ReactantBackend{Static.True, Missing, Nothing, AutoEnzyme{Nothing, Nothing}}(static(true), missing, nothing, AutoEnzyme()), static(false), (layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.014252576 -0.01007682 … -0.0066116317 -0.012908219; -0.0066900905 0.001967808 … -0.003597645 -0.0010663884; … ; -0.013411599 -0.009727624 … -0.0062034056 -0.012329484; -0.021803929 -0.008291057 … -0.010640304 -0.014436427]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-0.016153783, -0.0064898166, 0.0040475307, 0.048172455, 0.018123386, -0.026255043, 0.031807806, -0.00242692, 0.0047841636, 0.030119237 … -0.018572997, 0.002068058, 0.009790083, 0.018018095, -0.01424905, -0.060925536, -0.016411005, -0.014629741, -0.015240658, -0.023550106])), layer_2 = (layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-2.1789727f-5 -3.8227197f-5 … 2.2419797f-6 -6.324912f-5; 0.008564907 0.006512975 … -0.006969174 0.009865682; … ; -0.021595063 -0.021455316 … 0.013971797 -0.033741944; -0.023399103 -0.018417688 … 0.018593067 -0.028052688]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-8.567904f-5, 0.015442661, -0.012886962, 0.052984916, 0.000940611, -0.013805261, -0.025449416, 0.012329709, 0.09247952, -0.038467444 … 0.02124564, -0.012647957, -0.00045393303, -0.026561411, -0.025026277, -0.044307135, 0.033703856, 0.02700105, -0.049719024, -0.043526463])), layer_2 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.041988406 0.0035988672 … 0.012094588 -0.006939848; -0.06445931 0.02036321 … 0.044551518 0.008209803; … ; 0.24316937 -0.04692524 … -0.11571932 0.0070323865; 0.027664784 -0.008376195 … -0.018484466 -0.0030616056]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.04771936, 0.079492785, -0.063023046, 0.011547155, -0.4143036, -0.118838035, -0.16815118, -0.2632838, -0.2873201, -0.0339642])))), (compiled_gradient_function = Reactant compiled function compute_gradients_internal (with tag ##compute_gradients_internal_reactant#433), update_function = Reactant compiled function update! (with tag ##update!_reactant#686), compiled_grad_and_step_function = Reactant compiled function compute_gradients_internal_and_step! (with tag ##compute_gradients_internal_and_step!_reactant#1007), is_sharded = false)), GenericLossFunction{typeof(Lux.LossFunctionImpl.l2_distance_loss), typeof(Statistics.mean)}(Lux.LossFunctionImpl.l2_distance_loss, Statistics.mean), nothing, Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}}, Nothing}((layer_1 = Dense(128 => 256, tanh), layer_2 = Chain{@NamedTuple{layer_1::Dense{typeof(tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(256 => 256, tanh), layer_2 = Dense(256 => 10)), nothing)), nothing), (layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.2252784 0.22416925 … 0.1999151 -0.018334232; -0.022642436 0.15450513 … -0.06492576 0.18159999; … ; 0.038053643 -0.071250185 … -0.033061065 0.039139703; -0.18772855 -0.0965359 … -0.18063419 0.019627318]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.031100241, -0.059878908, 0.08487941, 2.0693342f-6, -0.06550352, -0.08487986, -0.026524307, 0.06386332, 0.042082705, 0.02731392 … -0.060524136, 0.0346705, -0.02837894, 0.06748107, 0.0027473217, -0.06903143, 0.006823757, 0.014109333, -0.029280972, 0.012123728])), layer_2 = (layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[0.120222874 -0.06340912 … -0.0006328311 -0.08487568; 0.05999857 0.08031667 … 0.10135459 -0.008547246; … ; -0.070574954 -0.12529366 … -0.112988465 0.03760839; 0.15816367 -0.12347525 … 0.01897169 0.1261452]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-0.047726065, -0.0076126982, 0.05891254, -0.05468241, 0.054435007, -0.005947113, -0.03260495, 0.025213236, 0.0024704752, 0.019756647 … -0.038096465, -0.029946823, 0.021725217, -0.004842341, 0.0018720245, -0.029938111, 0.008292873, 0.01365754, 0.053042553, -0.02033623])), layer_2 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.10451723 -0.027075663 … 0.018401911 0.06574812; -0.09850936 -0.10326829 … -0.021818167 0.047744956; … ; -0.061319478 -0.06494065 … 0.05045609 0.06725229; -0.025141204 0.094486676 … 0.02057817 -0.09107423]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-0.03402, 0.055178143, -0.036147784, 0.03546726, -0.026433855, -0.037797134, -0.02685109, -0.05179471, 0.05912261, -0.05172773])))), (layer_1 = NamedTuple(), layer_2 = (layer_1 = NamedTuple(), layer_2 = NamedTuple())), ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (layer_1 = (weight = Leaf(ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.00127828 -0.00108819 … -0.000579374 -0.00129518; -0.00129653 0.000363266 … -0.000695885 -0.000220152; … ; -0.00226619 -0.00174052 … -0.00104106 -0.00215551; -0.00363046 -0.00162021 … -0.00175398 -0.00258242]), Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[2.05798f-7 1.0234f-7 … 4.45381f-8 1.66622f-7; 9.3323f-8 7.29063f-9 … 2.68765f-8 2.72633f-9; … ; 2.85402f-7 1.67323f-7 … 6.03123f-8 2.56985f-7; 7.34734f-7 1.45927f-7 … 1.71924f-7 3.68345f-7]), (Reactant.ConcretePJRTNumber{Float32, 1}(0.729), Reactant.ConcretePJRTNumber{Float32, 1}(0.997003)))), bias = Leaf(ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-0.00147888, -0.00126067, -0.00024512, 0.0103003, 0.00511794, -0.00677713, 0.00692629, -0.000736556, 0.000429644, 0.00477295 … -0.00376831, 0.000957916, 0.000537737, 0.00326793, -0.00284112, -0.0107958, -0.00270653, -0.00281344, -0.00259105, -0.00396031]), Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[2.63239f-7, 8.82635f-8, 6.84696f-8, 6.02842f-6, 1.6761f-6, 2.81506f-6, 2.74192f-6, 3.59707f-8, 2.31813f-8, 1.28963f-6 … 7.95356f-7, 7.38564f-8, 1.1986f-7, 5.89749f-7, 4.50396f-7, 6.44006f-6, 4.09317f-7, 4.38955f-7, 3.72682f-7, 8.72425f-7]), (Reactant.ConcretePJRTNumber{Float32, 1}(0.729), Reactant.ConcretePJRTNumber{Float32, 1}(0.997003))))), layer_2 = (layer_1 = (weight = Leaf(ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-1.92071f-5 -1.98882f-6 … 2.45465f-5 9.11111f-8; 0.0017717 0.00124793 … -0.00150786 0.00187327; … ; -0.00373359 -0.0036783 … 0.00243024 -0.00579673; -0.00454264 -0.00359967 … 0.00358124 -0.0055097]), Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[3.62358f-11 1.87609f-12 … 7.29646f-11 9.0774f-12; 1.76661f-7 8.63204f-8 … 1.29674f-7 1.943f-7; … ; 7.71925f-7 7.50078f-7 … 3.2683f-7 1.8623f-6; 1.14592f-6 7.20328f-7 … 7.11382f-7 1.68898f-6]), (Reactant.ConcretePJRTNumber{Float32, 1}(0.729), Reactant.ConcretePJRTNumber{Float32, 1}(0.997003)))), bias = Leaf(ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-7.6764f-6, 0.002996, -0.00215631, 0.0108377, 0.000171212, -0.00324612, -0.00506657, 0.00160721, 0.0186094, -0.00646918 … 0.00417523, -0.0043, -0.000207296, -0.0046866, -0.00398426, -0.00686918, 0.00543797, 0.00308305, -0.00856162, -0.00854114]), Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[7.43883f-12, 4.98398f-7, 2.5891f-7, 6.59154f-6, 1.61884f-9, 6.19829f-7, 1.43188f-6, 1.69292f-7, 1.93607f-5, 2.32789f-6 … 9.70009f-7, 1.29615f-6, 3.43886f-9, 1.21397f-6, 8.97047f-7, 2.69644f-6, 1.66317f-6, 7.47133f-7, 4.06121f-6, 4.05819f-6]), (Reactant.ConcretePJRTNumber{Float32, 1}(0.729), Reactant.ConcretePJRTNumber{Float32, 1}(0.997003))))), layer_2 = (weight = Leaf(ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.00948004 0.00171439 … 0.00412935 -0.000683905; -0.00407593 0.0025052 … 0.00502547 0.00215895; … ; 0.0472378 -0.00898661 … -0.0214768 0.00263519; 0.0115309 -0.00133463 … -0.00363606 0.00186014]), Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[5.20286f-6 2.39227f-7 … 1.19777f-6 4.81734f-8; 4.84769f-6 4.41769f-7 … 2.02493f-6 2.88184f-7; … ; 0.000123925 4.47608f-6 … 2.54904f-5 5.09783f-7; 1.02391f-5 1.00626f-7 … 7.35784f-7 5.8815f-7]), (Reactant.ConcretePJRTNumber{Float32, 1}(0.729), Reactant.ConcretePJRTNumber{Float32, 1}(0.997003)))), bias = Leaf(ReactantOptimiser(Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.0001f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))), (Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.0111157, 0.00553539, -0.0061385, 0.00853476, -0.0791731, -0.0292157, -0.0252973, -0.0456424, -0.0556272, -0.0132202]), Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[7.24034f-6, 7.03766f-6, 3.97516f-6, 6.8506f-6, 0.000347333, 5.11703f-5, 3.71479f-5, 0.000115324, 0.000171764, 1.30559f-5]), (Reactant.ConcretePJRTNumber{Float32, 1}(0.729), Reactant.ConcretePJRTNumber{Float32, 1}(0.997003))))))), 2)) ``` ## Defining Custom Layers {#Defining-Custom-Layers} We can train our model using the above code, but let's go ahead and see how to use Reactant. Reactant is a julia frontend that generates MLIR and then compiles it using XLA (after running fancy optimizations). It is the current recommended way to train large models in Lux. For more details on using Reactant, see the [manual](/manual/compiling_lux_models#reactant-compilation). ```julia using Lux, Random, Optimisers, Reactant, Enzyme using Printf # For pretty printing dev = reactant_device() ``` ```ansi (::ReactantDevice{Missing, Missing, Missing, Missing, Union{}}) (generic function with 1 method) ``` We will define a custom MLP using the `@compact` macro. The macro takes in a list of parameters, layers and states, and a function defining the forward pass of the neural network. ```julia n_in = 1 n_out = 1 nlayers = 3 model = @compact( w1=Dense(n_in => 32), w2=[Dense(32 => 32) for i in 1:nlayers], w3=Dense(32 => n_out), act=relu ) do x embed = act(w1(x)) for w in w2 embed = act(w(embed)) end out = w3(embed) @return out end ``` ```ansi @compact( w1 = Dense(1 => 32), # 64 parameters w2 = NamedTuple( (1-3) = Dense(32 => 32), # 3_168 (1_056 x 3) parameters ), w3 = Dense(32 => 1), # 33 parameters act = relu, ) do x embed = act(w1(x)) for w = w2 embed = act(w(embed)) end out = w3(embed) return out end # Total: 3_265 parameters,  # plus 1 states. ``` We can initialize the model and train it with the same code as before! ```julia rng = Random.default_rng() Random.seed!(rng, 0) ps, st = Lux.setup(rng, model) |> dev x = rand(rng, Float32, n_in, 32) |> dev @jit model(x, ps, st) # 1×32 Matrix and updated state as output. x_data = reshape(collect(-2.0f0:0.1f0:2.0f0), 1, :) y_data = 2 .* x_data .- x_data .^ 3 x_data, y_data = dev(x_data), dev(y_data) function train_model!(model, ps, st, x_data, y_data, num_epochs=1000) train_state = Lux.Training.TrainState(model, ps, st, Adam(0.001f0)) for iter in 1:num_epochs _, loss, _, train_state = Lux.Training.single_train_step!( AutoEnzyme(), MSELoss(), (x_data, y_data), train_state ) if iter == 1 || iter % 100 == 0 || iter == num_epochs @printf "Iteration: %04d \t Loss: %10.9g\n" iter loss end end return model, train_state.parameters, train_state.states end train_model!(model, ps, st, x_data, y_data) ``` ```ansi Iteration: 0001 Loss: 2.08085132 Iteration: 0100 Loss: 0.13866809 Iteration: 0200 Loss: 0.00399366021 Iteration: 0300 Loss: 0.0010299189 Iteration: 0400 Loss: 0.000388183282 Iteration: 0500 Loss: 0.000311460026 Iteration: 0600 Loss: 0.000204949945 Iteration: 0700 Loss: 0.000141659068 Iteration: 0800 Loss: 0.000115743162 Iteration: 0900 Loss: 0.000106912601 Iteration: 1000 Loss: 0.000102958125 ``` ::: tip Training with Optimization.jl If you are coming from the SciML ecosystem and want to use Optimization.jl, please refer to the [Optimization.jl Tutorial](/tutorials/beginner/5_OptimizationIntegration#Optimization-Lux-Tutorial). ::: ## Additional Packages {#Additional-Packages} `LuxDL` hosts various packages that provide additional functionality for Lux.jl. All packages mentioned in this documentation are available via the Julia General Registry. You can install all those packages via `import Pkg; Pkg.add()`. ## Reactant & XLA (CPU/GPU/TPU) Support {#Reactant-and-XLA-CPU/GPU/TPU-Support} Lux.jl supports XLA compilation for CPU, GPU, and TPU using [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl). ## Native Julia GPU Support {#Native-Julia-GPU-Support} ::: warning Warning Using accelerators via Reactant and XLA is the preferred way to use GPUs with Lux.jl. ::: ::: danger AMD GPU Users For AMD GPUs, we **strongly recommend using Reactant** instead of native AMDGPU.jl support. Native AMDGPU.jl support is experimental with known issues including deadlocks. Reactant provides superior performance and reliability for AMD hardware. ::: GPU Support for Lux.jl requires loading additional packages: * [`LuxCUDA.jl`](https://github.com/LuxDL/LuxCUDA.jl) for CUDA support. * [`AMDGPU.jl`](https://github.com/JuliaGPU/AMDGPU.jl) for AMDGPU support (not recommended, use Reactant instead). * [`Metal.jl`](https://github.com/JuliaGPU/Metal.jl) for Apple Metal support. * [`oneAPI.jl`](https://github.com/JuliaGPU/oneAPI.jl) for oneAPI support. --- --- url: /dev/introduction/overview.md --- # Why we wrote Lux? {#Why-we-wrote-Lux?} Julia already has quite a few well established Neural Network Frameworks – [Flux](https://fluxml.ai/) & [KNet](https://denizyuret.github.io/Knet.jl/latest/). However, certain design elements – **Coupled Model and Parameters** & **Internal Mutations** – associated with these frameworks make them less compiler and user friendly. Making changes to address these problems in the respective frameworks would be too disruptive for users. Here comes in `Lux`: a neural network framework built completely using pure functions to make it both compiler and autodiff friendly. ## Performance and Deployment with Reactant {#Performance-and-Deployment-with-Reactant} Lux.jl takes a **Reactant-first approach** to deliver exceptional performance and seamless deployment capabilities: * **XLA Compilation** – Lux models compile to highly optimized XLA code via [Reactant.jl](https://github.com/EnzymeAD/Reactant.jl), delivering significant speedups on CPU, GPU, and TPU. * **Cross-Platform Performance** – Run the same Lux model with optimal performance across different hardware backends (CPU, NVIDIA GPUs, AMD GPUs, TPUs) without code changes, simply by switching the Reactant backend. * **Production Deployment** – Compiled models can be exported and deployed to production servers and edge devices by leveraging the rich TensorFlow ecosystem, making Lux suitable for real-world applications. * **Large Model Support** – With Reactant compilation, Lux now excels at training very large models that were previously challenging, making it competitive with other frameworks for large-scale deep learning. ## Design Principles {#Design-Principles} * **Layers must be immutable** – cannot store any parameter/state but rather store the information to construct them * **Layers are pure functions** * **Layers return a Tuple containing the result and the updated state** * **Given same inputs the outputs must be same** – yes this must hold true even for stochastic functions. Randomness must be controlled using `rng`s passed in the state. * **Easily extensible** * **Extensive Testing** – All layers and features are tested across all supported AD backends across all supported hardware backends. ## Why use Lux over Flux? {#Why-use-Lux-over-Flux?} * **High-Performance XLA Compilation** – Lux's Reactant-first approach enables XLA compilation for dramatic performance improvements across CPU, GPU, and TPU. Models compile to highly optimized code that eliminates Julia overhead and leverages hardware-specific optimizations. * **Production-Ready Deployment** – Deploy Lux models to production environments using the mature TensorFlow ecosystem. Compiled models can be exported and run on servers, edge devices, and mobile platforms. * **Neural Networks for SciML**: For SciML Applications (Neural ODEs, Deep Equilibrium Models) solvers typically expect a monolithic parameter vector. Flux enables this via its `destructure` mechanism, but `destructure` comes with various [edge cases and limitations](https://fluxml.ai/Optimisers.jl/dev/api/#Optimisers.destructure). Lux forces users to make an explicit distinction between state variables and parameter variables to avoid these issues. Also, it comes battery-included for distributed training. * **Sensible display of Custom Layers** – Ever wanted to see Pytorch like Network printouts or wondered how to extend the pretty printing of Flux's layers? Lux handles all of that by default. * **Truly immutable models** - No *unexpected internal mutations* since all layers are implemented as pure functions. All layers are also *deterministic* given the parameters and state: if a layer is supposed to be stochastic (say [`Lux.Dropout`](/api/Lux/layers#Lux.Dropout)), the state must contain a seed which is then updated after the function call. * **Easy Parameter Manipulation** – By separating parameter data and layer structures, Lux makes implementing [`WeightNorm`](/api/Lux/layers#Lux.WeightNorm), `SpectralNorm`, etc. downright trivial. Without this separation, it is much harder to pass such parameters around without mutations which AD systems don't like. * **Wider AD Support** – Lux has extensive support for most [AD systems in julia](/manual/autodiff#autodiff-lux), while Flux is mostly tied to Zygote (with some initial support for Enzyme). * **Optimized for All Model Sizes** – Whether you're working with small prototypes or large production models, Lux delivers optimal performance. For the smallest networks where minimal overhead is critical, you can use [`ToSimpleChainsAdaptor`](/api/Lux/interop#Lux.ToSimpleChainsAdaptor) to leverage SimpleChains.jl's specialized CPU optimizations. * **Reliability** – We have learned from the mistakes of the past with Flux and everything in our core framework is extensively tested, along with downstream CI to ensure that everything works as expected. --- --- url: /dev/introduction/resources.md --- # Resources to Get Started {#Resources-to-Get-Started} * Go through the [Quickstart Example](/introduction/index#Quickstart). * Read the introductory tutorials on [Julia](https://jump.dev/JuMP.jl/stable/tutorials/getting_started/getting_started_with_julia) and Lux. * Go through the examples sorted based on their complexity in the documentation. ::: tip Have More Questions? For usage related questions, please use [Github Discussions](https://github.com/orgs/LuxDL/discussions) which allows questions and answers to be indexed. To report bugs use [Github Issues](https://github.com/LuxDL/Lux.jl/issues) or even better send in a [Pull Request](https://github.com/LuxDL/Lux.jl/pulls). ::: --- --- url: /dev/introduction/updating_to_v1.md --- # Updating to Lux v1 {#updating-to-v1} Lux v1 is a Major Release, mostly to signify the stability of the API. In this page, we list out a concrete set of changes that need to be made to your code to update to Lux v1. We also list out some new exciting features that were added as part of this release. ## `LuxLib.jl` {#LuxLib.jl} ### Breaking Changes {#Breaking-Changes} * Old deprecated API with keyword arguments has been removed. See the new docs in [LuxLib API](/api/NN_Primitives/LuxLib#LuxLib-API) for more details. * Default for [`layernorm`](/api/NN_Primitives/LuxLib#LuxLib.API.layernorm) dims has been changed to exclude the batch dimension. ### New Major Features {#New-Major-Features} * Dense layers now support CUDA backend for Enzyme (starting `v1.1`). Wider support for other operations with Enzyme + CUDA is being actively worked on. ## `LuxCore.jl` {#LuxCore.jl} ### Breaking Changes {#Breaking-Changes-2} * `AbstractExplicitLayer` has been renamed to `AbstractLuxLayer`. * `AbstractExplicitContainerLayer` behaviour * This has been renamed to `AbstractLuxContainerLayer`. * Previously, `AbstractExplicitContainerLayer{(:a,)}` (i.e. singleton containers) would produce default initial parameters and states without wrapping them in a `NamedTuple{(:a,)}`. This was inconsistent with non-singleton containers, and was a source of confusion. With `v` we return `(; a = )` and `(; a = )` by default. See [`AbstractLuxWrapperLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxWrapperLayer) for a replacement of this functionality. * `inputsize` has been removed since it was ambiguous and not used anywhere. * Changes to `outputsize`: * Single argument version has been removed. See [LuxCore.jl Pull Request 43](https://github.com/LuxDL/LuxCore.jl/pull/43#issuecomment-2254232817) for more details on the rationale behind this change. * Fallback implementation has been moved to `Lux.jl`. (i.e. users using Lux shouldn't see a difference, but if `Lux.jl` isn't loaded, this function has error.) * Internally this uses a `NilArray` that is able to compute sizes without actually running the computation. * `Functors` and `Setfield` have been made into optional dependencies. Certain `LuxCore` functionality that rely on these functions, will throw an error if these packages are not loaded. ### New Major Features {#New-Major-Features-2} * Introduction of [`AbstractLuxWrapperLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxWrapperLayer). This behaves exactly like the old singleton container. For example, the old `AbstractExplicitContainerLayer{(:a,)}` is equivalent to `AbstractLuxWrapperLayer{:a}`. ## `WeightInitializers.jl` {#WeightInitializers.jl} This was a major release to signify the stability of the API. There were no breaking changes. We do support a wider range of RNG types, see [Supported RNG Types](/api/Building_Blocks/WeightInitializers#Supported-RNG-Types-WeightInit) for more details. ## `MLDataDevices.jl` {#MLDataDevices.jl} This is the most aggressive change that was made. We renamed the `LuxDeviceUtils.jl` package to `MLDataDevices.jl`, to allow for non-Lux packages to use this shared device management abstraction. ::: warning Deprecation of `LuxDeviceUtils.jl` This also marks the deprecation of the `LuxDeviceUtils.jl` package. We won't be making any updates to that package, including fixing any bugs. All users should switch to `MLDataDevices.jl` instead. ::: ### Breaking Changes {#Breaking-Changes-3} * `Lux(___)Device` objects have been renamed to `(___)Device`. For example, `LuxCUDADevice` has been renamed to `CUDADevice`. * `Lux(___)Adaptor` objects have been removed. The corresponding `Device` objects should be used directly instead. ### New Major Features {#New-Major-Features-3} * [`DeviceIterator`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.DeviceIterator) provides a generalization of `CUDA.CuIterator` and works for all backends and more data types (using `Functors.jl`). `MLUtils.DataLoader |> gdev` now returns a `DeviceIterator` instead of being a no-op. ## `Lux.jl` {#Lux.jl} ### Breaking Changes (Removed Functionality) {#Breaking-Changes-Removed-Functionality} * Direct reexport of `NNlib` has been removed. We reexport selected functionality from `NNlib`. Direactly load `NNlib` if you need to use the other functions. * Flattening of [`Chain`](/api/Lux/layers#Lux.Chain) layers has been removed, and the corresponding `disable_optimizations` kwarg has been removed. * Some layers overloaded `Base.keys`, these have been removed. These were mostly un-documented and weren't supposed to be used outside of the `Lux.jl` package. * [`Training.TrainState`](/api/Lux/utilities#Lux.Training.TrainState) construction with `rng` has been removed. * Older versions of Preferences have been removed. * `disable_stacktrace_truncation!` has been removed. From Julia 1.9 onwards, stacktrace truncation is enabled by default. * Certain Experimental features were present outside the `Lux.Experimental` module. These have been removed, use them via `Lux.Experimental` instead. Run Julia with with `depwarn` as `error` and Lux `v0.5` to see the deprecations. * `Lux.Experimental.@layer_map` is not longer needed and has been removed. The name of the variable prevents writing generic functions and is no longer pre-pended to the `KeyPath`. See the docstring of [`Lux.Experimental.layer_map`](/api/Lux/contrib#Lux.Experimental.layer_map) for more details. * `allow_fast_activation` kwarg has been removed completely. Pass an anonymous function as the activation to prevent internal modivations to the activation function. ### Breaking Changes (Moved Functionality) {#Breaking-Changes-Moved-Functionality} * `Lux.Experimental.Training` has been moved to `Lux.Training`. We guarantee SemVar on this new module. * `Lux.cpu` and `Lux.gpu` have been removed. Use [`cpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.cpu_device) and [`gpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_device) instead. * `Experimental.@compact` can be directly used via [`@compact`](/api/Lux/utilities#Lux.@compact) now. * `Experimental.StatefulLuxLayer` has been moved to [`Lux.StatefulLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer). * `st_fixed_path` kwarg has been removed from [`Lux.StatefulLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer), instead use it as `StatefulLuxLayer{st_fixed_path}(...)`. * Strings as inputs to [`Lux.Experimental.layer_map`](/api/Lux/contrib#Lux.Experimental.layer_map) and [`Lux.Experimental.@debug_mode`](/api/Lux/contrib#Lux.Experimental.@debug_mode) are removed, use `Functors.KeyPath` instead. * `CrossCor` has been removed. Use `Conv(args...; kwargs..., cross_correlation=true)` instead. ### Breaking Changes (Changes in Defaults) {#Breaking-Changes-Changes-in-Defaults} * [`Conv`](/api/Lux/layers#Lux.Conv) and [`ConvTranspose`](/api/Lux/layers#Lux.ConvTranspose) use an initialization based on the activation function, taken from Pytorch. Pytorch assumes the activation function is `leakyrelu` to compute the gain, however, we compute the gain based on the activation function passed in to the layer. * [`Upsample`](/api/Lux/layers#Lux.Upsample) now has an `align_corners` keyword argument, which defaults to `false`. Previously this was always `true`. * [`Dense`](/api/Lux/layers#Lux.Dense) and [`Bilinear`](/api/Lux/layers#Lux.Bilinear) have updated default initializations to align with the defaults from Pytorch. See the documentation for more details. * [`InstanceNorm`](/api/Lux/layers#Lux.InstanceNorm) now defaults to `affine=false` instead of `affine=true`. * [`Embedding`](/api/Lux/layers#Lux.Embedding) now defaults to `init_weight=rand32` instead of `init_weight=randn32`. * Recurrent Cells - [`RNNCell`](/api/Lux/layers#Lux.RNNCell), [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell), and [`GRUCell`](/api/Lux/layers#Lux.GRUCell) now have different default initializations. See the documentation for more details. ### New Features {#New-Features} * [`InstanceNorm`](/api/Lux/layers#Lux.InstanceNorm) now supports tracking statistics. * [`RNNCell`](/api/Lux/layers#Lux.RNNCell) and [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell) add `bias_ih` and `bias_hh` to the parameters to align with Pytorch. Both are controlled using `init_bias` and `use_bias`. * [`ConvTranspose`](/api/Lux/layers#Lux.ConvTranspose) allows `flipkernel=true` via `cross_correlation=true`. This makes it efficient for MIOpen. * [`ConvTranspose`](/api/Lux/layers#Lux.ConvTranspose) now has an `outpad` keyword argument, which is used to increase the size of the output in the desired dimensions. * Pooling Layers based on lpnorm have been added – [`LPPool`](/api/Lux/layers#Lux.LPPool), [`GlobalLPPool`](/api/Lux/layers#Lux.GlobalLPPool), and [`AdaptiveLPPool`](/api/Lux/layers#Lux.AdaptiveLPPool). --- --- url: /dev/introduction/citation.md --- # Citation {#Citation} If you found this library to be useful in academic work, then please cite: ```bibtex @software{pal2023lux, author = {Pal, Avik}, title = , month = {April}, year = 2023, note = {If you use this software, please cite it as below.}, publisher = {Zenodo}, version = {v0.5.0}, doi = {10.5281/zenodo.7808904}, url = {https://doi.org/10.5281/zenodo.7808904} } ``` ```bibtex @thesis{pal2023efficient, title = , author = {Pal, Avik}, year = {2023}, school = {Massachusetts Institute of Technology} } ``` --- --- url: /dev/tutorials.md --- # Tutorials ## Beginner Tutorials ## Intermediate Tutorials ## Advanced Tutorials ## Selected 3rd Party Tutorials ::: warning These tutorials are developed by the community and may not be up-to-date with the latest version of `Lux.jl`. Please refer to the official documentation for the most up-to-date information. Please open an issue (ideally both at `Lux.jl` and at the downstream linked package) if any of them are non-functional and we will try to get them updated. ::: ::: tip If you found an amazing tutorial showcasing `Lux.jl` online, or wrote one yourself, please open an issue or PR to add it to the list! ::: --- --- url: /dev/tutorials/beginner/1_Basics.md --- # Julia & Lux for the Uninitiated {#Julia-and-Lux-for-the-Uninitiated} This is a quick intro to [Lux](https://github.com/LuxDL/Lux.jl) loosely based on: 1. [PyTorch's tutorial](https://pytorch.org/tutorials/beginner/deep_learning_60min_blitz.html). 2. Flux's tutorial (the link for which has now been lost to abyss). 3. [Jax's tutorial](https://jax.readthedocs.io/en/latest/jax-101/index.html). It introduces basic Julia programming, as well `Zygote`, a source-to-source automatic differentiation (AD) framework in Julia. We'll use these tools to build a very simple neural network. Let's start with importing `Lux.jl` ```julia using Lux, Random ``` ``` Precompiling packages... 423.3 ms ✓ StructUtilsStaticArraysCoreExt (serial) 1 dependency successfully precompiled in 0 seconds ``` Now let us control the randomness in our code using proper Pseudo Random Number Generator (PRNG) ```julia rng = Random.default_rng() Random.seed!(rng, 0) ``` ``` Random.TaskLocalRNG() ``` ## Arrays {#Arrays} The starting point for all of our models is the `Array` (sometimes referred to as a `Tensor` in other frameworks). This is really just a list of numbers, which might be arranged into a shape like a square. Let's write down an array with three elements. ```julia x = [1, 2, 3] ``` ``` 3-element Vector{Int64}: 1 2 3 ``` Here's a matrix – a square array with four elements. ```julia x = [1 2; 3 4] ``` ``` 2×2 Matrix{Int64}: 1 2 3 4 ``` We often work with arrays of thousands of elements, and don't usually write them down by hand. Here's how we can create an array of 5×3 = 15 elements, each a random number from zero to one. ```julia x = rand(rng, 5, 3) ``` ``` 5×3 Matrix{Float64}: 0.455238 0.746943 0.193291 0.547642 0.746801 0.116989 0.773354 0.97667 0.899766 0.940585 0.0869468 0.422918 0.0296477 0.351491 0.707534 ``` There's a few functions like this; try replacing `rand` with `ones`, `zeros`, or `randn`. By default, Julia works stores numbers is a high-precision format called `Float64`. In ML we often don't need all those digits, and can ask Julia to work with `Float32` instead. We can even ask for more digits using `BigFloat`. ```julia x = rand(BigFloat, 5, 3) ``` ``` 5×3 Matrix{BigFloat}: 0.981339 0.793159 0.459019 0.043883 0.624384 0.56055 0.164786 0.524008 0.0355555 0.414769 0.577181 0.621958 0.00823197 0.30215 0.655881 ``` ```julia x = rand(Float32, 5, 3) ``` ``` 5×3 Matrix{Float32}: 0.567794 0.369178 0.342539 0.0985227 0.201145 0.587206 0.776598 0.148248 0.0851708 0.723731 0.0770206 0.839303 0.404728 0.230954 0.679087 ``` We can ask the array how many elements it has. ```julia length(x) ``` ``` 15 ``` Or, more specifically, what size it has. ```julia size(x) ``` ``` (5, 3) ``` We sometimes want to see some elements of the array on their own. ```julia x ``` ``` 5×3 Matrix{Float32}: 0.567794 0.369178 0.342539 0.0985227 0.201145 0.587206 0.776598 0.148248 0.0851708 0.723731 0.0770206 0.839303 0.404728 0.230954 0.679087 ``` ```julia x[2, 3] ``` ``` 0.58720636f0 ``` This means get the second row and the third column. We can also get every row of the third column. ```julia x[:, 3] ``` ``` 5-element Vector{Float32}: 0.34253937 0.58720636 0.085170805 0.8393034 0.67908657 ``` We can add arrays, and subtract them, which adds or subtracts each element of the array. ```julia x + x ``` ``` 5×3 Matrix{Float32}: 1.13559 0.738356 0.685079 0.197045 0.40229 1.17441 1.5532 0.296496 0.170342 1.44746 0.154041 1.67861 0.809456 0.461908 1.35817 ``` ```julia x - x ``` ``` 5×3 Matrix{Float32}: 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ``` Julia supports a feature called *broadcasting*, using the `.` syntax. This tiles small arrays (or single numbers) to fill bigger ones. ```julia x .+ 1 ``` ``` 5×3 Matrix{Float32}: 1.56779 1.36918 1.34254 1.09852 1.20114 1.58721 1.7766 1.14825 1.08517 1.72373 1.07702 1.8393 1.40473 1.23095 1.67909 ``` We can see Julia tile the column vector `1:5` across all rows of the larger array. ```julia zeros(5, 5) .+ (1:5) ``` ``` 5×5 Matrix{Float64}: 1.0 1.0 1.0 1.0 1.0 2.0 2.0 2.0 2.0 2.0 3.0 3.0 3.0 3.0 3.0 4.0 4.0 4.0 4.0 4.0 5.0 5.0 5.0 5.0 5.0 ``` The x' syntax is used to transpose a column `1:5` into an equivalent row, and Julia will tile that across columns. ```julia zeros(5, 5) .+ (1:5)' ``` ``` 5×5 Matrix{Float64}: 1.0 2.0 3.0 4.0 5.0 1.0 2.0 3.0 4.0 5.0 1.0 2.0 3.0 4.0 5.0 1.0 2.0 3.0 4.0 5.0 1.0 2.0 3.0 4.0 5.0 ``` We can use this to make a times table. ```julia (1:5) .* (1:5)' ``` ``` 5×5 Matrix{Int64}: 1 2 3 4 5 2 4 6 8 10 3 6 9 12 15 4 8 12 16 20 5 10 15 20 25 ``` Finally, and importantly for machine learning, we can conveniently do things like matrix multiply. ```julia W = randn(5, 10) x = rand(10) W * x ``` ``` 5-element Vector{Float64}: 1.2197981041108443 -2.62625877100596 -2.8573820474674845 -2.4319346874291314 1.0108668577150213 ``` Julia's arrays are very powerful, and you can learn more about what they can do [here](https://docs.julialang.org/en/v1/manual/arrays/). ### CUDA Arrays {#CUDA-Arrays} CUDA functionality is provided separately by the [CUDA.jl package](https://github.com/JuliaGPU/CUDA.jl). If you have a GPU and LuxCUDA is installed, Lux will provide CUDA capabilities. For additional details on backends see the manual section. You can manually add `CUDA`. Once CUDA is loaded you can move any array to the GPU with the `cu` function (or the `gpu` function exported by `Lux`), and it supports all of the above operations with the same syntax. ```julia using LuxCUDA if LuxCUDA.functional() x_cu = cu(rand(5, 3)) @show x_cu end ``` ## (Im)mutability {#Immutability} Lux as you might have read is "Immutable by convention," which means that the core library is built without any form of mutation and all functions are pure. However, we don't enforce it in any form. We do **strongly recommend** that users extending this framework for their respective applications don't mutate their arrays. ```julia x = reshape(1:8, 2, 4) ``` ``` 2×4 reshape(::UnitRange{Int64}, 2, 4) with eltype Int64: 1 3 5 7 2 4 6 8 ``` To update this array, we should first copy the array. ```julia x_copy = copy(x) view(x_copy, :, 1) .= 0 println("Original Array ", x) println("Mutated Array ", x_copy) ``` ``` Original Array [1 3 5 7; 2 4 6 8] Mutated Array [0 3 5 7; 0 4 6 8] ``` Note that our current default AD engine (Zygote) is unable to differentiate through this mutation, however, for these specialized cases it is quite trivial to write custom backward passes. (This problem will be fixed once we move towards Enzyme.jl) ## Managing Randomness {#Managing-Randomness} We rely on the Julia StdLib `Random` for managing the randomness in our execution. First, we create an PRNG (pseudorandom number generator) and seed it. ```julia rng = Xoshiro(0) # Creates a Xoshiro PRNG with seed 0 ``` ``` Random.Xoshiro(0xdb2fa90498613fdf, 0x48d73dc42d195740, 0x8c49bc52dc8a77ea, 0x1911b814c02405e8, 0x22a21880af5dc689) ``` If we call any function that relies on `rng` and uses it via `randn`, `rand`, etc. `rng` will be mutated. As we have already established we care a lot about immutability, hence we should use `Lux.replicate` on PRNGs before using them. First, let us run a random number generator 3 times with the `replicate`d rng. ```julia random_vectors = Vector{Vector{Float64}}(undef, 3) for i in 1:3 random_vectors[i] = rand(Lux.replicate(rng), 10) println("Iteration $i ", random_vectors[i]) end @assert random_vectors[1] ≈ random_vectors[2] ≈ random_vectors[3] ``` ``` Iteration 1 [0.4552384158732863, 0.5476424498276177, 0.7733535276924052, 0.9405848223512736, 0.02964765308691042, 0.74694291453392, 0.7468008914093891, 0.9766699015845924, 0.08694684883050086, 0.35149138733595564] Iteration 2 [0.4552384158732863, 0.5476424498276177, 0.7733535276924052, 0.9405848223512736, 0.02964765308691042, 0.74694291453392, 0.7468008914093891, 0.9766699015845924, 0.08694684883050086, 0.35149138733595564] Iteration 3 [0.4552384158732863, 0.5476424498276177, 0.7733535276924052, 0.9405848223512736, 0.02964765308691042, 0.74694291453392, 0.7468008914093891, 0.9766699015845924, 0.08694684883050086, 0.35149138733595564] ``` As expected we get the same output. We can remove the `replicate` call and we will get different outputs. ```julia for i in 1:3 println("Iteration $i ", rand(rng, 10)) end ``` ``` Iteration 1 [0.4552384158732863, 0.5476424498276177, 0.7733535276924052, 0.9405848223512736, 0.02964765308691042, 0.74694291453392, 0.7468008914093891, 0.9766699015845924, 0.08694684883050086, 0.35149138733595564] Iteration 2 [0.018743665453639813, 0.8601828553599953, 0.6556360448565952, 0.7746656838366666, 0.7817315740767116, 0.5553797706980106, 0.1261990389976131, 0.4488101521328277, 0.624383955429775, 0.05657739601024536] Iteration 3 [0.19597391412112541, 0.6830945313415872, 0.6776220912718907, 0.6456416023530093, 0.6340362477836592, 0.5595843665394066, 0.5675557670686644, 0.34351700231383653, 0.7237308297251812, 0.3691778381831775] ``` ## Automatic Differentiation {#Automatic-Differentiation} Julia has quite a few (maybe too many) AD tools. For the purpose of this tutorial, we will use: 1. [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) – For Jacobian-Vector Product (JVP) 2. [Zygote.jl](https://github.com/FluxML/Zygote.jl) – For Vector-Jacobian Product (VJP) *Slight Detour*: We have had several questions regarding if we will be considering any other AD system for the reverse-diff backend. For now we will stick to Zygote.jl, however once we have tested Lux extensively with [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl), we will make the switch. Even though, theoretically, a VJP (Vector-Jacobian product - reverse autodiff) and a JVP (Jacobian-Vector product - forward-mode autodiff) are similar—they compute a product of a Jacobian and a vector—they differ by the computational complexity of the operation. In short, when you have a large number of parameters (hence a wide matrix), a JVP is less efficient computationally than a VJP, and, conversely, a JVP is more efficient when the Jacobian matrix is a tall matrix. ```julia using ComponentArrays, ForwardDiff, Zygote ``` ``` Precompiling packages... 432.2 ms ✓ StructUtilsTablesExt (serial) 1 dependency successfully precompiled in 0 seconds ``` ### Gradients {#Gradients} For our first example, consider a simple function computing $f(x) = \frac{1}{2}x^T x$, where $\nabla f(x) = x$ ```julia f(x) = x' * x / 2 ∇f(x) = x # `∇` can be typed as `\nabla` v = randn(rng, Float32, 4) ``` ``` 4-element Vector{Float32}: -0.4051151 -0.4593922 0.92155594 1.1871622 ``` Let's use ForwardDiff and Zygote to compute the gradients. ```julia println("Actual Gradient: ", ∇f(v)) println("Computed Gradient via Reverse Mode AD (Zygote): ", only(Zygote.gradient(f, v))) println("Computed Gradient via Forward Mode AD (ForwardDiff): ", ForwardDiff.gradient(f, v)) ``` ``` Actual Gradient: Float32[-0.4051151, -0.4593922, 0.92155594, 1.1871622] Computed Gradient via Reverse Mode AD (Zygote): Float32[-0.4051151, -0.4593922, 0.92155594, 1.1871622] Computed Gradient via Forward Mode AD (ForwardDiff): Float32[-0.4051151, -0.4593922, 0.92155594, 1.1871622] ``` Note that `AD.gradient` will only work for scalar valued outputs. ### Jacobian-Vector Product {#Jacobian-Vector-Product} I will defer the discussion on forward-mode AD to . Here let us just look at a mini example on how to use it. ```julia f(x) = x .* x ./ 2 x = randn(rng, Float32, 5) v = ones(Float32, 5) ``` ``` 5-element Vector{Float32}: 1.0 1.0 1.0 1.0 1.0 ``` ::: warning Using DifferentiationInterface While DifferentiationInterface provides these functions for a wider range of backends, we currently don't recommend using them with Lux models, since the functions presented here come with additional goodies like [fast second-order derivatives](/manual/nested_autodiff#nested_autodiff). ::: Compute the JVP. `AutoForwardDiff` specifies that we want to use `ForwardDiff.jl` for the Jacobian-Vector Product ```julia jvp = jacobian_vector_product(f, AutoForwardDiff(), x, v) println("JVP: ", jvp) ``` ``` JVP: Float32[-0.877497, 1.1953009, -0.057005208, 0.25055695, 0.09351656] ``` ### Vector-Jacobian Product {#Vector-Jacobian-Product} Using the same function and inputs, let us compute the Vector-Jacobian Product (VJP). ```julia vjp = vector_jacobian_product(f, AutoZygote(), x, v) println("VJP: ", vjp) ``` ``` VJP: Float32[-0.877497, 1.1953009, -0.057005208, 0.25055695, 0.09351656] ``` ## Linear Regression {#Linear-Regression} Finally, now let us consider a linear regression problem. From a set of data-points ${ (x\_i, y\_i), i \in { 1, \dots, k }, x\_i \in \mathbb{R}^n, y\_i \in \mathbb{R}^m }$, we try to find a set of parameters $W$ and $b$, such that $f\_{W,b}(x) = Wx + b$, which minimizes the mean squared error: $$L(W, b) \longrightarrow \sum\_{i = 1}^{k} \frac{1}{2} | y\_i - f\_{W,b}(x\_i) |\_2^2$$ We can write `f` from scratch, but to demonstrate `Lux`, let us use the `Dense` layer. ```julia model = Dense(10 => 5) rng = Random.default_rng() Random.seed!(rng, 0) ``` ``` Random.TaskLocalRNG() ``` Let us initialize the parameters and states (in this case it is empty) for the model. ```julia ps, st = Lux.setup(rng, model) ps = ComponentArray(ps) ``` ``` ComponentVector{Float32}(weight = Float32[-0.48351604 0.29944378 0.44048923 0.52216566 0.20001544 0.14378412 4.831728f-6 0.53108513 -0.30674055 0.034259237; -0.049033877 -0.42427677 0.27051237 0.40789896 -0.43846488 -0.17706363 -0.032581452 0.46514037 0.19584312 0.23992884; 0.4501613 0.48263645 -0.29908532 -0.18695378 -0.110237636 -0.44184566 0.40354213 0.25278288 0.18056089 -0.35231933; 0.05218965 -0.09701933 0.27035677 0.12589002 -0.2956183 0.34717596 -0.421895 -0.1307367 0.36829442 -0.30972943; 0.20277861 -0.5152452 -0.22635894 0.18841727 0.29828638 0.21690919 -0.04265763 -0.41919124 0.07148273 -0.45247707], bias = Float32[-0.04199602, -0.093925126, -0.0007736237, -0.19397983, 0.0066712513]) ``` Set problem dimensions. ```julia n_samples = 20 x_dim = 10 y_dim = 5 ``` We're going to generate a random set of weights `W` and biases `b` that will act as our true model (also known as the ground truth). The neural network we'll train will be to try and approximate `W` and `b` from example data. ```julia W = randn(rng, Float32, y_dim, x_dim) b = randn(rng, Float32, y_dim) ``` Generate samples with additional noise. ```julia x_samples = randn(rng, Float32, x_dim, n_samples) y_samples = W * x_samples .+ b .+ 0.01f0 .* randn(rng, Float32, y_dim, n_samples) println("x shape: ", size(x_samples), "; y shape: ", size(y_samples)) ``` ``` x shape: (10, 20); y shape: (5, 20) ``` For updating our parameters let's use [Optimisers.jl](https://github.com/FluxML/Optimisers.jl). We will use Stochastic Gradient Descent (SGD) with a learning rate of `0.01`. ```julia using Optimisers, Printf ``` Define the loss function ```julia lossfn = MSELoss() println("Loss Value with ground true parameters: ", lossfn(W * x_samples .+ b, y_samples)) ``` ``` Loss Value with ground true parameters: 9.3742405e-5 ``` We will train the model using our training API. ```julia function train_model!(model, ps, st, opt, nepochs::Int) tstate = Training.TrainState(model, ps, st, opt) for i in 1:nepochs grads, loss, _, tstate = Training.single_train_step!( AutoZygote(), lossfn, (x_samples, y_samples), tstate ) if i == 1 || i % 1000 == 0 || i == nepochs @printf "Loss Value after %6d iterations: %.8f\n" i loss end end return tstate.model, tstate.parameters, tstate.states end model, ps, st = train_model!(model, ps, st, Descent(0.01f0), 10000) println("Loss Value after training: ", lossfn(first(model(x_samples, ps, st)), y_samples)) ``` ``` Loss Value after 1 iterations: 7.80465460 Loss Value after 1000 iterations: 0.12503763 Loss Value after 2000 iterations: 0.02538623 Loss Value after 3000 iterations: 0.00914946 Loss Value after 4000 iterations: 0.00407888 Loss Value after 5000 iterations: 0.00198553 Loss Value after 6000 iterations: 0.00101213 Loss Value after 7000 iterations: 0.00053365 Loss Value after 8000 iterations: 0.00029220 Loss Value after 9000 iterations: 0.00016886 Loss Value after 10000 iterations: 0.00010551 Loss Value after training: 0.00010546902 ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/beginner/2_PolynomialFitting.md --- # Fitting a Polynomial using MLP {#Fitting-a-Polynomial-using-MLP} In this tutorial we will fit a MultiLayer Perceptron (MLP) on data generated from a polynomial. ## Package Imports {#Package-Imports} ```julia using Lux, ADTypes, Optimisers, Printf, Random, Reactant, Statistics, CairoMakie ``` ## Dataset {#Dataset} Generate 128 datapoints from the polynomial $y = x^2 - 2x$. ```julia function generate_data(rng::AbstractRNG) x = reshape(collect(range(-2.0f0, 2.0f0, 128)), (1, 128)) poly_coeffs = (0, -2, 1) y = evalpoly.(x, (poly_coeffs,)) # add some noise to simulate real-world conditions y .+= randn(rng, Float32, (1, 128)) .* 0.1f0 return (x, y) end ``` Initialize the random number generator and fetch the dataset. ```julia rng = MersenneTwister() Random.seed!(rng, 12345) (x, y) = generate_data(rng) ``` ``` (Float32[-2.0 -1.968504 -1.9370079 -1.9055119 -1.8740157 -1.8425196 -1.8110236 -1.7795275 -1.7480315 -1.7165354 -1.6850394 -1.6535434 -1.6220472 -1.5905511 -1.5590551 -1.527559 -1.496063 -1.464567 -1.4330709 -1.4015749 -1.3700787 -1.3385826 -1.3070866 -1.2755905 -1.2440945 -1.2125984 -1.1811024 -1.1496063 -1.1181102 -1.0866141 -1.0551181 -1.023622 -0.992126 -0.96062994 -0.92913383 -0.8976378 -0.86614174 -0.8346457 -0.8031496 -0.77165353 -0.7401575 -0.70866144 -0.6771653 -0.6456693 -0.61417323 -0.5826772 -0.5511811 -0.51968503 -0.48818898 -0.4566929 -0.42519686 -0.39370078 -0.36220473 -0.33070865 -0.2992126 -0.26771653 -0.23622048 -0.20472442 -0.17322835 -0.14173229 -0.11023622 -0.07874016 -0.047244094 -0.015748031 0.015748031 0.047244094 0.07874016 0.11023622 0.14173229 0.17322835 0.20472442 0.23622048 0.26771653 0.2992126 0.33070865 0.36220473 0.39370078 0.42519686 0.4566929 0.48818898 0.51968503 0.5511811 0.5826772 0.61417323 0.6456693 0.6771653 0.70866144 0.7401575 0.77165353 0.8031496 0.8346457 0.86614174 0.8976378 0.92913383 0.96062994 0.992126 1.023622 1.0551181 1.0866141 1.1181102 1.1496063 1.1811024 1.2125984 1.2440945 1.2755905 1.3070866 1.3385826 1.3700787 1.4015749 1.4330709 1.464567 1.496063 1.527559 1.5590551 1.5905511 1.6220472 1.6535434 1.6850394 1.7165354 1.7480315 1.7795275 1.8110236 1.8425196 1.8740157 1.9055119 1.9370079 1.968504 2.0], Float32[8.080871 7.562357 7.451749 7.5005703 7.295229 7.2245107 6.8731666 6.7092047 6.5385857 6.4631066 6.281978 5.960991 5.963052 5.68927 5.3667717 5.519665 5.2999034 5.0238676 5.174298 4.6706038 4.570324 4.439068 4.4462147 4.299262 3.9799082 3.9492173 3.8747025 3.7264304 3.3844414 3.2934628 3.1180353 3.0698316 3.0491123 2.592982 2.8164148 2.3875027 2.3781595 2.4269633 2.2763796 2.3316176 2.0829067 1.9049499 1.8581494 1.7632381 1.7745113 1.5406592 1.3689325 1.2614254 1.1482575 1.2801026 0.9070533 0.91188717 0.9415703 0.85747254 0.6692604 0.7172643 0.48259094 0.48990166 0.35299227 0.31578436 0.25483933 0.37486005 0.19847682 -0.042415008 -0.05951088 0.014774345 -0.114184186 -0.15978265 -0.29916334 -0.22005874 -0.17161606 -0.3613516 -0.5489093 -0.7267406 -0.5943626 -0.62129945 -0.50063384 -0.6346849 -0.86081326 -0.58715504 -0.5171875 -0.6575044 -0.71243864 -0.78395927 -0.90537953 -0.9515314 -0.8603811 -0.92880917 -1.0078154 -0.90215015 -1.0109437 -1.0764086 -1.1691734 -1.0740278 -1.1429857 -1.104191 -0.948015 -0.9233653 -0.82379496 -0.9810639 -0.92863405 -0.9360056 -0.92652786 -0.847396 -1.115507 -1.0877254 -0.92295444 -0.86975616 -0.81879705 -0.8482455 -0.6524158 -0.6184501 -0.7483137 -0.60395515 -0.67555165 -0.6288941 -0.6774449 -0.49889082 -0.43817532 -0.46497717 -0.30316323 -0.36745527 -0.3227286 -0.20977046 -0.09777648 -0.053120755 -0.15877295 -0.06777584]) ``` Let's visualize the dataset ```julia begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y") l = lines!( ax, x[1, :], x -> evalpoly(x, (0, -2, 1)); linewidth=3, color=:blue, label="True Quadratic Function", ) s = scatter!( ax, x[1, :], y[1, :]; markersize=12, alpha=0.5, color=:orange, strokecolor=:black, strokewidth=2, label="Actual Data", ) axislegend(ax) fig end ``` ## Neural Network {#Neural-Network} For this problem, you should not be using a neural network. But let's still do that! ```julia model = Chain(Dense(1 => 16, relu), Dense(16 => 1)) ``` ``` Chain( layer_1 = Dense(1 => 16, relu), # 32 parameters layer_2 = Dense(16 => 1), # 17 parameters ) # Total: 49 parameters, # plus 0 states. ``` ## Optimizer {#Optimizer} We will use Adam from [Optimisers.jl](https://fluxml.ai/Optimisers.jl) ```julia opt = Adam(0.03f0) ``` ``` Optimisers.Adam(eta=0.03, beta=(0.9, 0.999), epsilon=1.0e-8) ``` ## Loss Function {#Loss-Function} We will use the `Training` API so we need to ensure that our loss function takes 4 inputs – model, parameters, states and data. The function must return 3 values – loss, updated\_state, and any computed statistics. This is already satisfied by the loss functions provided by Lux. ```julia const loss_function = MSELoss() const cdev = cpu_device() const xdev = reactant_device() ps, st = Lux.setup(rng, model) |> xdev ``` ``` ((layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[2.2569513; 1.8385266; 1.8834435; -1.4215803; -0.1289033; -1.4116536; -1.4359436; -2.3610642; -0.847535; 1.6091344; -0.34999675; 1.9372884; -0.41628727; 1.1786895; -1.4312565; 0.34652048;;]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.9155488, -0.005158901, 0.5026965, -0.84174657, -0.9167142, -0.14881086, -0.8202727, 0.19286752, 0.60171676, 0.951689, 0.4595859, -0.33281517, -0.692657, 0.4369135, 0.3800323, 0.61768365])), layer_2 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[0.20061705 0.22529833 0.07667785 0.115506485 0.22827768 0.22680467 0.0035893882 -0.39495495 0.18033011 -0.02850357 -0.08613788 -0.3103005 0.12508307 -0.087390475 -0.13759731 0.08034529]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.06066203]))), (layer_1 = NamedTuple(), layer_2 = NamedTuple())) ``` ## Training {#Training} First we will create a [`Training.TrainState`](/api/Lux/utilities#Lux.Training.TrainState) which is essentially a convenience wrapper over parameters, states and optimizer states. ```julia tstate = Training.TrainState(model, ps, st, opt) ``` ``` TrainState( Chain( layer_1 = Dense(1 => 16, relu), # 32 parameters layer_2 = Dense(16 => 1), # 17 parameters ), number of parameters: 49 number of states: 0 optimizer: ReactantOptimiser(Optimisers.Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.03f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))) step: 0 ) ``` Now we will use Enzyme (Reactant) for our AD requirements. ```julia vjp_rule = AutoEnzyme() ``` Finally the training loop. ```julia function main(tstate::Training.TrainState, vjp, data, epochs) data = xdev(data) for epoch in 1:epochs _, loss, _, tstate = Training.single_train_step!(vjp, loss_function, data, tstate) if epoch % 50 == 1 || epoch == epochs @printf "Epoch: %3d \t Loss: %.5g\n" epoch loss end end return tstate end tstate = main(tstate, vjp_rule, (x, y), 250) ``` ``` TrainState( Chain( layer_1 = Dense(1 => 16, relu), # 32 parameters layer_2 = Dense(16 => 1), # 17 parameters ), number of parameters: 49 number of states: 0 optimizer: ReactantOptimiser(Optimisers.Adam(eta=Reactant.ConcretePJRTNumber{Float32, 1}(0.03f0), beta=(Reactant.ConcretePJRTNumber{Float64, 1}(0.9), Reactant.ConcretePJRTNumber{Float64, 1}(0.999)), epsilon=Reactant.ConcretePJRTNumber{Float64, 1}(1.0e-8))) step: 250 cache: TrainingBackendCache(Lux.Training.ReactantBackend{Static.True, Missing, Nothing, ADTypes.AutoEnzyme{Nothing, Nothing}}(static(true), missing, nothing, ADTypes.AutoEnzyme())) objective_function: GenericLossFunction ) ``` Since we are using Reactant, we need to compile the model before we can use it. ```julia forward_pass = @compile Lux.apply( tstate.model, xdev(x), tstate.parameters, Lux.testmode(tstate.states) ) y_pred = forward_pass(tstate.model, xdev(x), tstate.parameters, Lux.testmode(tstate.states)) |> first |> cdev ``` Let's plot the results ```julia begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y") l = lines!( ax, x[1, :], x -> evalpoly(x, (0, -2, 1)); linewidth=3, label="True Quadratic Function", ) s1 = scatter!( ax, x[1, :], y[1, :]; markersize=12, alpha=0.5, color=:orange, strokecolor=:black, strokewidth=2, label="Actual Data", ) s2 = scatter!( ax, x[1, :], y_pred[1, :]; markersize=12, alpha=0.5, color=:green, strokecolor=:black, strokewidth=2, label="Predictions", ) axislegend(ax) fig end ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/beginner/3_SimpleRNN.md --- # Training a Simple LSTM {#Training-a-Simple-LSTM} In this tutorial we will go over using a recurrent neural network to classify clockwise and anticlockwise spirals. By the end of this tutorial you will be able to: 1. Create custom Lux models. 2. Become familiar with the Lux recurrent neural network API. 3. Training using Optimisers.jl and Zygote.jl. ## Package Imports {#Package-Imports} Note: If you wish to use `AutoZygote()` for automatic differentiation, add Zygote to your project dependencies and include `using Zygote`. ```julia using ADTypes, Lux, JLD2, MLUtils, Optimisers, Printf, Reactant, Random ``` ## Dataset {#Dataset} We will use MLUtils to generate 500 (noisy) clockwise and 500 (noisy) anticlockwise spirals. Using this data we will create a `MLUtils.DataLoader`. Our dataloader will give us sequences of size 2 × seq\_len × batch\_size and we need to predict a binary value whether the sequence is clockwise or anticlockwise. ```julia function create_dataset(; dataset_size=1000, sequence_length=50) # Create the spirals data = [MLUtils.Datasets.make_spiral(sequence_length) for _ in 1:dataset_size] # Get the labels labels = vcat(repeat([0.0f0], dataset_size ÷ 2), repeat([1.0f0], dataset_size ÷ 2)) clockwise_spirals = [ reshape(d[1][:, 1:sequence_length], :, sequence_length, 1) for d in data[1:(dataset_size ÷ 2)] ] anticlockwise_spirals = [ reshape(d[1][:, (sequence_length + 1):end], :, sequence_length, 1) for d in data[((dataset_size ÷ 2) + 1):end] ] x_data = Float32.(cat(clockwise_spirals..., anticlockwise_spirals...; dims=3)) return x_data, labels end function get_dataloaders(; dataset_size=1000, sequence_length=50) x_data, labels = create_dataset(; dataset_size, sequence_length) # Split the dataset (x_train, y_train), (x_val, y_val) = splitobs((x_data, labels); at=0.8, shuffle=true) # Create DataLoaders return ( # Use DataLoader to automatically minibatch and shuffle the data DataLoader( collect.((x_train, y_train)); batchsize=128, shuffle=true, partial=false ), # Don't shuffle the validation data DataLoader(collect.((x_val, y_val)); batchsize=128, shuffle=false, partial=false), ) end ``` ## Creating a Classifier {#Creating-a-Classifier} We will be extending the `Lux.AbstractLuxContainerLayer` type for our custom model since it will contain a LSTM block and a classifier head. We pass the field names `lstm_cell` and `classifier` to the type to ensure that the parameters and states are automatically populated and we don't have to define `Lux.initialparameters` and `Lux.initialstates`. To understand more about container layers, please look at [Container Layer](/manual/interface#Container-Layer). ```julia struct SpiralClassifier{L,C} <: AbstractLuxContainerLayer{(:lstm_cell, :classifier)} lstm_cell::L classifier::C end ``` We won't define the model from scratch but rather use the [`Lux.LSTMCell`](/api/Lux/layers#Lux.LSTMCell) and [`Lux.Dense`](/api/Lux/layers#Lux.Dense). ```julia function SpiralClassifier(in_dims, hidden_dims, out_dims) return SpiralClassifier( LSTMCell(in_dims => hidden_dims), Dense(hidden_dims => out_dims, sigmoid) ) end ``` We can use default Lux blocks – `Recurrence(LSTMCell(in_dims => hidden_dims)` – instead of defining the following. But let's still do it for the sake of it. Now we need to define the behavior of the Classifier when it is invoked. ```julia function (s::SpiralClassifier)( x::AbstractArray{T,3}, ps::NamedTuple, st::NamedTuple ) where {T} # First we will have to run the sequence through the LSTM Cell # The first call to LSTM Cell will create the initial hidden state # See that the parameters and states are automatically populated into a field called # `lstm_cell` We use `eachslice` to get the elements in the sequence without copying, # and `Iterators.peel` to split out the first element for LSTM initialization. x_init, x_rest = Iterators.peel(LuxOps.eachslice(x, Val(2))) (y, carry), st_lstm = s.lstm_cell(x_init, ps.lstm_cell, st.lstm_cell) # Now that we have the hidden state and memory in `carry` we will pass the input and # `carry` jointly for x in x_rest (y, carry), st_lstm = s.lstm_cell((x, carry), ps.lstm_cell, st_lstm) end # After running through the sequence we will pass the output through the classifier y, st_classifier = s.classifier(y, ps.classifier, st.classifier) # Finally remember to create the updated state st = merge(st, (classifier=st_classifier, lstm_cell=st_lstm)) return vec(y), st end ``` ## Using the `@compact` API {#Using-the-@compact-API} We can also define the model using the [`Lux.@compact`](/api/Lux/utilities#Lux.@compact) API, which is a more concise way of defining models. This macro automatically handles the boilerplate code for you and as such we recommend this way of defining custom layers ```julia function SpiralClassifierCompact(in_dims, hidden_dims, out_dims) return @compact(; lstm_cell=LSTMCell(in_dims => hidden_dims), classifier=Dense(hidden_dims => out_dims, sigmoid) ) do x::AbstractArray{T,3} where {T} x_init, x_rest = Iterators.peel(LuxOps.eachslice(x, Val(2))) y, carry = lstm_cell(x_init) for x in x_rest y, carry = lstm_cell((x, carry)) end @return vec(classifier(y)) end end ``` ## Defining Accuracy, Loss and Optimiser {#Defining-Accuracy,-Loss-and-Optimiser} Now let's define the binary cross-entropy loss. Typically it is recommended to use `logitbinarycrossentropy` since it is more numerically stable, but for the sake of simplicity we will use `binarycrossentropy`. ```julia const lossfn = BinaryCrossEntropyLoss() function compute_loss(model, ps, st, (x, y)) ŷ, st_ = model(x, ps, st) loss = lossfn(ŷ, y) return loss, st_, (; y_pred=ŷ) end matches(y_pred, y_true) = sum((y_pred .> 0.5f0) .== y_true) accuracy(y_pred, y_true) = matches(y_pred, y_true) / length(y_pred) ``` ## Training the Model {#Training-the-Model} ```julia function main(model_type) dev = reactant_device() cdev = cpu_device() # Get the dataloaders train_loader, val_loader = get_dataloaders() |> dev # Create the model model = model_type(2, 8, 1) ps, st = Lux.setup(Random.default_rng(), model) |> dev train_state = Training.TrainState(model, ps, st, Adam(0.01f0)) model_compiled = if dev isa ReactantDevice @compile model(first(train_loader)[1], ps, Lux.testmode(st)) else model end ad = dev isa ReactantDevice ? AutoReactant() : AutoZygote() for epoch in 1:25 # Train the model total_loss = 0.0f0 total_samples = 0 for (x, y) in train_loader (_, loss, _, train_state) = Training.single_train_step!( ad, lossfn, (x, y), train_state ) total_loss += loss * length(y) total_samples += length(y) end @printf("Epoch [%3d]: Loss %4.5f\n", epoch, total_loss / total_samples) # Validate the model total_acc = 0.0f0 total_loss = 0.0f0 total_samples = 0 st_ = Lux.testmode(train_state.states) for (x, y) in val_loader ŷ, st_ = model_compiled(x, train_state.parameters, st_) ŷ, y = cdev(ŷ), cdev(y) total_acc += accuracy(ŷ, y) * length(y) total_loss += lossfn(ŷ, y) * length(y) total_samples += length(y) end @printf( "Validation:\tLoss %4.5f\tAccuracy %4.5f\n", total_loss / total_samples, total_acc / total_samples ) end return (train_state.parameters, train_state.states) |> cdev end ps_trained, st_trained = main(SpiralClassifier) ``` ``` ┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`. └ @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/LuxCore.jl:18 Epoch [ 1]: Loss 0.66542 Validation: Loss 0.58600 Accuracy 1.00000 Epoch [ 2]: Loss 0.54916 Validation: Loss 0.48532 Accuracy 1.00000 Epoch [ 3]: Loss 0.45214 Validation: Loss 0.39267 Accuracy 1.00000 Epoch [ 4]: Loss 0.36055 Validation: Loss 0.30246 Accuracy 1.00000 Epoch [ 5]: Loss 0.27154 Validation: Loss 0.21883 Accuracy 1.00000 Epoch [ 6]: Loss 0.19208 Validation: Loss 0.14930 Accuracy 1.00000 Epoch [ 7]: Loss 0.12941 Validation: Loss 0.09960 Accuracy 1.00000 Epoch [ 8]: Loss 0.08676 Validation: Loss 0.06817 Accuracy 1.00000 Epoch [ 9]: Loss 0.06044 Validation: Loss 0.04918 Accuracy 1.00000 Epoch [ 10]: Loss 0.04483 Validation: Loss 0.03803 Accuracy 1.00000 Epoch [ 11]: Loss 0.03531 Validation: Loss 0.03081 Accuracy 1.00000 Epoch [ 12]: Loss 0.02910 Validation: Loss 0.02584 Accuracy 1.00000 Epoch [ 13]: Loss 0.02467 Validation: Loss 0.02229 Accuracy 1.00000 Epoch [ 14]: Loss 0.02140 Validation: Loss 0.01956 Accuracy 1.00000 Epoch [ 15]: Loss 0.01898 Validation: Loss 0.01741 Accuracy 1.00000 Epoch [ 16]: Loss 0.01690 Validation: Loss 0.01566 Accuracy 1.00000 Epoch [ 17]: Loss 0.01528 Validation: Loss 0.01421 Accuracy 1.00000 Epoch [ 18]: Loss 0.01388 Validation: Loss 0.01298 Accuracy 1.00000 Epoch [ 19]: Loss 0.01269 Validation: Loss 0.01194 Accuracy 1.00000 Epoch [ 20]: Loss 0.01172 Validation: Loss 0.01102 Accuracy 1.00000 Epoch [ 21]: Loss 0.01084 Validation: Loss 0.01022 Accuracy 1.00000 Epoch [ 22]: Loss 0.01004 Validation: Loss 0.00950 Accuracy 1.00000 Epoch [ 23]: Loss 0.00934 Validation: Loss 0.00885 Accuracy 1.00000 Epoch [ 24]: Loss 0.00868 Validation: Loss 0.00826 Accuracy 1.00000 Epoch [ 25]: Loss 0.00812 Validation: Loss 0.00773 Accuracy 1.00000 ``` We can also train the compact model with the exact same code! ```julia ps_trained2, st_trained2 = main(SpiralClassifierCompact) ``` ``` ┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`. └ @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/LuxCore.jl:18 Epoch [ 1]: Loss 0.79162 Validation: Loss 0.70687 Accuracy 0.52344 Epoch [ 2]: Loss 0.66215 Validation: Loss 0.59054 Accuracy 0.90625 Epoch [ 3]: Loss 0.55249 Validation: Loss 0.48464 Accuracy 1.00000 Epoch [ 4]: Loss 0.44685 Validation: Loss 0.38201 Accuracy 1.00000 Epoch [ 5]: Loss 0.34736 Validation: Loss 0.29338 Accuracy 1.00000 Epoch [ 6]: Loss 0.26724 Validation: Loss 0.22457 Accuracy 1.00000 Epoch [ 7]: Loss 0.20342 Validation: Loss 0.16954 Accuracy 1.00000 Epoch [ 8]: Loss 0.15195 Validation: Loss 0.12521 Accuracy 1.00000 Epoch [ 9]: Loss 0.11194 Validation: Loss 0.08959 Accuracy 1.00000 Epoch [ 10]: Loss 0.08033 Validation: Loss 0.06608 Accuracy 1.00000 Epoch [ 11]: Loss 0.05962 Validation: Loss 0.05055 Accuracy 1.00000 Epoch [ 12]: Loss 0.04616 Validation: Loss 0.04007 Accuracy 1.00000 Epoch [ 13]: Loss 0.03708 Validation: Loss 0.03296 Accuracy 1.00000 Epoch [ 14]: Loss 0.03075 Validation: Loss 0.02761 Accuracy 1.00000 Epoch [ 15]: Loss 0.02585 Validation: Loss 0.02325 Accuracy 1.00000 Epoch [ 16]: Loss 0.02159 Validation: Loss 0.01931 Accuracy 1.00000 Epoch [ 17]: Loss 0.01796 Validation: Loss 0.01637 Accuracy 1.00000 Epoch [ 18]: Loss 0.01548 Validation: Loss 0.01444 Accuracy 1.00000 Epoch [ 19]: Loss 0.01384 Validation: Loss 0.01311 Accuracy 1.00000 Epoch [ 20]: Loss 0.01265 Validation: Loss 0.01204 Accuracy 1.00000 Epoch [ 21]: Loss 0.01167 Validation: Loss 0.01114 Accuracy 1.00000 Epoch [ 22]: Loss 0.01083 Validation: Loss 0.01037 Accuracy 1.00000 Epoch [ 23]: Loss 0.01010 Validation: Loss 0.00968 Accuracy 1.00000 Epoch [ 24]: Loss 0.00945 Validation: Loss 0.00907 Accuracy 1.00000 Epoch [ 25]: Loss 0.00886 Validation: Loss 0.00851 Accuracy 1.00000 ``` ## Saving the Model {#Saving-the-Model} We can save the model using JLD2 (and any other serialization library of your choice) Note that we transfer the model to CPU before saving. Additionally, we recommend that you don't save the model struct and only save the parameters and states. ```julia @save "trained_model.jld2" ps_trained st_trained ``` Let's try loading the model ```julia @load "trained_model.jld2" ps_trained st_trained ``` ``` 2-element Vector{Symbol}: :ps_trained :st_trained ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, icelake-server) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/beginner/4_SimpleChains.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # MNIST Classification with SimpleChains {#MNIST-Classification-with-SimpleChains} SimpleChains.jl is an excellent framework for training small neural networks. In this tutorial we will demonstrate how to use the same API as Lux.jl to train a model using SimpleChains.jl. We will use the tutorial from [SimpleChains.jl](https://pumasai.github.io/SimpleChains.jl/dev/examples/mnist/) as a reference. ## Package Imports {#Package-Imports} ```julia using Lux, MLUtils, Optimisers, Zygote, OneHotArrays, Random, Statistics, Printf, Reactant using MLDatasets: MNIST using SimpleChains: SimpleChains Reactant.set_default_backend("cpu") ``` ## Loading MNIST {#Loading-MNIST} ```julia function loadmnist(batchsize, train_split) # Load MNIST N = parse(Bool, get(ENV, "CI", "false")) ? 1500 : nothing dataset = MNIST(; split=:train) if N !== nothing imgs = dataset.features[:, :, 1:N] labels_raw = dataset.targets[1:N] else imgs = dataset.features labels_raw = dataset.targets end # Process images into (H, W, C, BS) batches x_data = Float32.(reshape(imgs, size(imgs, 1), size(imgs, 2), 1, size(imgs, 3))) y_data = onehotbatch(labels_raw, 0:9) (x_train, y_train), (x_test, y_test) = splitobs((x_data, y_data); at=train_split) return ( # Use DataLoader to automatically minibatch and shuffle the data DataLoader(collect.((x_train, y_train)); batchsize, shuffle=true, partial=false), # Don't shuffle the test data DataLoader(collect.((x_test, y_test)); batchsize, shuffle=false, partial=false), ) end ``` ## Define the Model {#Define-the-Model} ```julia lux_model = Chain( Conv((5, 5), 1 => 6, relu), MaxPool((2, 2)), Conv((5, 5), 6 => 16, relu), MaxPool((2, 2)), FlattenLayer(3), Chain(Dense(256 => 128, relu), Dense(128 => 84, relu), Dense(84 => 10)), ) ``` We now need to convert the lux\_model to SimpleChains.jl. We need to do this by defining the `ToSimpleChainsAdaptor` and providing the input dimensions. ```julia adaptor = ToSimpleChainsAdaptor((28, 28, 1)) simple_chains_model = adaptor(lux_model) ``` ## Helper Functions {#Helper-Functions} ```julia const lossfn = CrossEntropyLoss(; logits=Val(true)) function accuracy(model, ps, st, dataloader) total_correct, total = 0, 0 st = Lux.testmode(st) for (x, y) in dataloader target_class = onecold(Array(y)) predicted_class = onecold(Array(first(model(x, ps, st)))) total_correct += sum(target_class .== predicted_class) total += length(target_class) end return total_correct / total end ``` ## Define the Training Loop {#Define-the-Training-Loop} ```julia function train(model, dev=cpu_device(); rng=Random.default_rng(), kwargs...) train_dataloader, test_dataloader = loadmnist(128, 0.9) |> dev ps, st = Lux.setup(rng, model) |> dev vjp = dev isa ReactantDevice ? AutoEnzyme() : AutoZygote() train_state = Training.TrainState(model, ps, st, Adam(3.0f-4)) if dev isa ReactantDevice x_ra = first(test_dataloader)[1] model_compiled = @compile model(x_ra, ps, Lux.testmode(st)) else model_compiled = model end ### Lets train the model nepochs = 10 tr_acc, te_acc = 0.0, 0.0 for epoch in 1:nepochs stime = time() for (x, y) in train_dataloader _, _, _, train_state = Training.single_train_step!( vjp, lossfn, (x, y), train_state ) end ttime = time() - stime tr_acc = accuracy( model_compiled, train_state.parameters, train_state.states, train_dataloader ) * 100 te_acc = accuracy( model_compiled, train_state.parameters, train_state.states, test_dataloader ) * 100 @printf "[%2d/%2d] \t Time %.2fs \t Training Accuracy: %.2f%% \t Test Accuracy: \ %.2f%%\n" epoch nepochs ttime tr_acc te_acc end return tr_acc, te_acc end ``` ## Finally Training the Model {#Finally-Training-the-Model} First we will train the Lux model ```julia tr_acc, te_acc = train(lux_model, reactant_device()) ``` Now we will train the SimpleChains model ```julia tr_acc, te_acc = train(simple_chains_model) ``` On my local machine we see a 3-4x speedup when using SimpleChains.jl. The conditions of the server this documentation is being built on is not ideal for CPU benchmarking hence, the speedup may not be as significant and even there might be regressions. *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/beginner/5_OptimizationIntegration.md --- # Fitting with Optimization.jl {#Optimization-Lux-Tutorial} Lux's native [Training.TrainState](/api/Lux/utilities#Lux.Training.TrainState) is a great API for gradient-based learning of neural networks, however, it is geared towards using `Optimisers.jl` as the backend. However, often times we want to train the neural networks with other optimization methods like BFGS, LBFGS, etc. In this tutorial, we will show how to train Lux models with Optimization.jl that provides a simple unified interface to various optimization methods. We will base our tutorial on the minibatching tutorial from the official [Optimization.jl](https://docs.sciml.ai/Optimization/stable/tutorials/minibatch/) docs. ::: tip Neural ODE This tutorial uses a Neural ODE, however, we won't discuss that part in this tutorial. Please refer to the Neural ODE tutorial for more information. ::: ## Imports packages {#Imports-packages} ```julia using Lux, Optimization, OptimizationOptimisers, OptimizationOptimJL, OrdinaryDiffEqTsit5, SciMLSensitivity, Random, MLUtils, CairoMakie, ComponentArrays, Printf const gdev = gpu_device() const cdev = cpu_device() ``` ``` ┌ Warning: No functional GPU backend found! Defaulting to CPU. │ │ 1. If no GPU is available, nothing needs to be done. Set `MLDATADEVICES_SILENCE_WARN_NO_GPU=1` to silence this warning. │ 2. If GPU is available, load the corresponding trigger package. │ a. `CUDA.jl` and `cuDNN.jl` (or just `LuxCUDA.jl`) for NVIDIA CUDA Support. │ b. `AMDGPU.jl` for AMD GPU ROCM Support. │ c. `Metal.jl` for Apple Metal GPU Support. (Experimental) │ d. `oneAPI.jl` for Intel oneAPI GPU Support. (Experimental) │ e. `OpenCL.jl` for OpenCL support. (Experimental) └ @ MLDataDevices.Internal ~/work/Lux.jl/Lux.jl/lib/MLDataDevices/src/internal.jl:114 ``` ## Generate some training data {#Generate-some-training-data} ```julia function lotka_volterra(du, u, p, t) x, y = u α, β, δ, γ = p du[1] = α * x - β * x * y du[2] = -δ * y + γ * x * y return nothing end u0 = [1.0f0, 1.0f0] datasize = 32 tspan = (0.0f0, 2.0f0) const t = range(tspan[1], tspan[2]; length=datasize) true_prob = ODEProblem(lotka_volterra, u0, (tspan[1], tspan[2]), [1.5, 1.0, 3.0, 1.0]) const ode_data = Array(solve(true_prob, Tsit5(); saveat=t)) begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]) lines!(ax, t, ode_data[1, :]; label=L"u_1(t)", color=:blue, linestyle=:dot, linewidth=4) lines!(ax, t, ode_data[2, :]; label=L"u_2(t)", color=:red, linestyle=:dot, linewidth=4) axislegend(ax; position=:lt) fig end ``` ## Define the DataLoader {#Define-the-DataLoader} We will define the DataLoader to batch over the data, additionally we will pipe it through the `gdev` device to move the data to the GPU on each iteration. By default `gdev` will move all objects to the GPU. But we don't want to move the time vector to the GPU. So we will wrap it in a struct and mark it as a leaf using MLDataDevices.isleaf ```julia struct TimeWrapper{T} t::T end MLDataDevices.isleaf(::TimeWrapper) = true Base.length(t::TimeWrapper) = length(t.t) Base.getindex(t::TimeWrapper, i) = TimeWrapper(t.t[i]) dataloader = DataLoader((ode_data, TimeWrapper(t)); batchsize=8) |> gdev ``` ## Training the model {#Training-the-model} Here we are using different optimization methods for demonstration purposes. This problem is trivial enough to not require this. Optimization.jl requires an abstract array as the parameters, hence we will construct a `ComponentArray` to store the parameters. ::: tip Parameter Estimation vs State Estimation Optimization.jl performs state estimation, which effectively means for a function `f(u, p)`, it is trying to compute the optimal `u` for a given `p`. This terminology might be confusing to ML practitioners, since in the ML world, we usually do parameter estimation. This effectively means that the `u` in Optimization.jl corresponds to our model parameters that is being optimized. ::: ```julia function train_model(dataloader) model = Chain(Dense(2, 32, tanh), Dense(32, 32, tanh), Dense(32, 2)) ps, st = Lux.setup(Random.default_rng(), model) ps_ca = ComponentArray(ps) |> gdev st = st |> gdev function callback(state, l) if state.iter == 1 || state.iter % 25 == 0 @printf "Iteration: %5d, Loss: %.6e\n" state.iter l end return l < 1.0e-8 ## Terminate if loss is small end smodel = StatefulLuxLayer(model, nothing, st) function loss_adjoint(θ, (u_batch, t_batch)) t_batch = t_batch.t u0 = u_batch[:, 1] dudt(u, p, t) = smodel(u, p) prob = ODEProblem(dudt, u0, (t_batch[1], t_batch[end]), θ) sol = solve(prob, Tsit5(); sensealg=InterpolatingAdjoint(), saveat=t_batch) pred = stack(sol.u) return MSELoss()(pred, u_batch) end # Define the Optimization Function that takes in the optimization state (our parameters) # and optimization parameters (nothing in our case) and data from the dataloader and # returns the loss. opt_func = OptimizationFunction(loss_adjoint, Optimization.AutoZygote()) opt_prob = OptimizationProblem(opt_func, ps_ca, dataloader) epochs = 25 res_adam = solve(opt_prob, Optimisers.Adam(0.001); callback, epochs) # Let's finetune a bit with L-BFGS opt_prob = OptimizationProblem(opt_func, res_adam.u, (gdev(ode_data), TimeWrapper(t))) res_lbfgs = solve(opt_prob, LBFGS(); callback, maxiters=epochs) # Now that we have a good fit, let's train it on the entire dataset without # Minibatching. We need to do this since ODE solves can lead to accumulated errors if # the model was trained on individual parts (without a data-shooting approach). opt_prob = remake(opt_prob; u0=res_lbfgs.u) res = solve(opt_prob, Optimisers.Adam(0.005); maxiters=500, callback) return StatefulLuxLayer(model, res.u, smodel.st) end trained_model = train_model(dataloader) ``` ``` Iteration: 1, Loss: 1.740769e-01 Iteration: 25, Loss: 8.644896e-02 Iteration: 50, Loss: 5.306665e-02 Iteration: 75, Loss: 8.427274e-02 Iteration: 100, Loss: 1.756596e-01 Iteration: 100, Loss: 2.055030e-02 Iteration: 1, Loss: 3.068537e-01 Iteration: 25, Loss: 3.350112e-02 Iteration: 1, Loss: 3.302021e-02 Iteration: 25, Loss: 4.320877e-02 Iteration: 50, Loss: 3.087694e-02 Iteration: 75, Loss: 2.725422e-02 Iteration: 100, Loss: 2.353250e-02 Iteration: 125, Loss: 2.144463e-02 Iteration: 150, Loss: 1.901578e-02 Iteration: 175, Loss: 1.681522e-02 Iteration: 200, Loss: 2.023219e-02 Iteration: 225, Loss: 1.649002e-02 Iteration: 250, Loss: 1.468544e-02 Iteration: 275, Loss: 1.324297e-02 Iteration: 300, Loss: 2.042319e-02 Iteration: 325, Loss: 1.605902e-02 Iteration: 350, Loss: 1.312578e-02 Iteration: 375, Loss: 1.188488e-02 Iteration: 400, Loss: 1.076655e-02 Iteration: 425, Loss: 1.006581e-02 Iteration: 450, Loss: 1.139796e-02 Iteration: 475, Loss: 9.892777e-03 Iteration: 500, Loss: 8.870113e-03 Iteration: 500, Loss: 8.870113e-03 ``` ## Plotting the results {#Plotting-the-results} ```julia dudt(u, p, t) = trained_model(u, p) prob = ODEProblem(dudt, gdev(u0), (tspan[1], tspan[2]), trained_model.ps) sol = solve(prob, Tsit5(); saveat=t) pred = convert(AbstractArray, sol) |> cdev begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]) lines!(ax, t, ode_data[1, :]; label=L"u_1(t)", color=:blue, linestyle=:dot, linewidth=4) lines!(ax, t, ode_data[2, :]; label=L"u_2(t)", color=:red, linestyle=:dot, linewidth=4) lines!(ax, t, pred[1, :]; label=L"\hat{u}_1(t)", color=:blue, linewidth=4) lines!(ax, t, pred[2, :]; label=L"\hat{u}_2(t)", color=:red, linewidth=4) axislegend(ax; position=:lt) fig end ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/1_NeuralODE.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # MNIST Classification using Neural ODEs {#MNIST-Classification-using-Neural-ODEs} To understand Neural ODEs, users should look up [these lecture notes](https://book.sciml.ai/notes/11-Differentiable_Programming_and_Neural_Differential_Equations/). We recommend users to directly use [DiffEqFlux.jl](https://docs.sciml.ai/DiffEqFlux/stable/), instead of implementing Neural ODEs from scratch. ## Package Imports {#Package-Imports} ```julia using Lux, ComponentArrays, SciMLSensitivity, LuxCUDA, Optimisers, OrdinaryDiffEqTsit5, Random, Statistics, Zygote, OneHotArrays, InteractiveUtils, Printf using MLDatasets: MNIST using MLUtils: DataLoader, splitobs CUDA.allowscalar(false) ``` ## Loading MNIST {#Loading-MNIST} ```julia function loadmnist(batchsize, train_split) # Load MNIST: Only 1500 for demonstration purposes N = parse(Bool, get(ENV, "CI", "false")) ? 1500 : nothing dataset = MNIST(; split=:train) if N !== nothing imgs = dataset.features[:, :, 1:N] labels_raw = dataset.targets[1:N] else imgs = dataset.features labels_raw = dataset.targets end # Process images into (H,W,C,BS) batches x_data = Float32.(reshape(imgs, size(imgs, 1), size(imgs, 2), 1, size(imgs, 3))) y_data = onehotbatch(labels_raw, 0:9) (x_train, y_train), (x_test, y_test) = splitobs((x_data, y_data); at=train_split) return ( # Use DataLoader to automatically minibatch and shuffle the data DataLoader(collect.((x_train, y_train)); batchsize, shuffle=true), # Don't shuffle the test data DataLoader(collect.((x_test, y_test)); batchsize, shuffle=false), ) end ``` ## Define the Neural ODE Layer {#Define-the-Neural-ODE-Layer} First we will use the `@compact` macro to define the Neural ODE Layer. ```julia function NeuralODECompact( model::Lux.AbstractLuxLayer; solver=Tsit5(), tspan=(0.0f0, 1.0f0), kwargs... ) return @compact(; model, solver, tspan, kwargs...) do x, p dudt(u, p, t) = vec(model(reshape(u, size(x)), p)) # Note the `p.model` here prob = ODEProblem(ODEFunction{false}(dudt), vec(x), tspan, p.model) @return solve(prob, solver; kwargs...) end end ``` We recommend using the compact macro for creating custom layers. The below implementation exists mostly for historical reasons when `@compact` was not part of the stable API. Also, it helps users understand how the layer interface of Lux works. The NeuralODE is a ContainerLayer, which stores a `model`. The parameters and states of the NeuralODE are same as those of the underlying model. ```julia struct NeuralODE{M<:Lux.AbstractLuxLayer,So,T,K} <: Lux.AbstractLuxWrapperLayer{:model} model::M solver::So tspan::T kwargs::K end function NeuralODE( model::Lux.AbstractLuxLayer; solver=Tsit5(), tspan=(0.0f0, 1.0f0), kwargs... ) return NeuralODE(model, solver, tspan, kwargs) end ``` [OrdinaryDiffEq.jl](https://docs.sciml.ai/OrdinaryDiffEq/stable/) can deal with non-Vector Inputs! However, certain discrete sensitivities like `ReverseDiffAdjoint` can't handle non-Vector inputs. Hence, we need to convert the input and output of the ODE solver to a Vector. ```julia function (n::NeuralODE)(x, ps, st) function dudt(u, p, t) u_, st = n.model(reshape(u, size(x)), p, st) return vec(u_) end prob = ODEProblem{false}(ODEFunction{false}(dudt), vec(x), n.tspan, ps) return solve(prob, n.solver; n.kwargs...), st end @views diffeqsol_to_array(l::Int, x::ODESolution) = reshape(last(x.u), (l, :)) @views diffeqsol_to_array(l::Int, x::AbstractMatrix) = reshape(x[:, end], (l, :)) ``` ## Create and Initialize the Neural ODE Layer {#Create-and-Initialize-the-Neural-ODE-Layer} ```julia function create_model( model_fn=NeuralODE; dev=gpu_device(), use_named_tuple::Bool=false, sensealg=InterpolatingAdjoint(; autojacvec=ZygoteVJP()), ) # Construct the Neural ODE Model model = Chain( FlattenLayer(), Dense(784 => 20, tanh), model_fn( Chain(Dense(20 => 10, tanh), Dense(10 => 10, tanh), Dense(10 => 20, tanh)); save_everystep=false, reltol=1.0f-3, abstol=1.0f-3, save_start=false, sensealg, ), Base.Fix1(diffeqsol_to_array, 20), Dense(20 => 10), ) rng = Random.default_rng() Random.seed!(rng, 0) ps, st = Lux.setup(rng, model) if !use_named_tuple ps = ComponentArray(ps) end ps = ps |> dev st = st |> dev return model, ps, st end ``` ## Define Utility Functions {#Define-Utility-Functions} ```julia const logitcrossentropy = CrossEntropyLoss(; logits=Val(true)) function accuracy(model, ps, st, dataloader) total_correct, total = 0, 0 st = Lux.testmode(st) for (x, y) in dataloader target_class = onecold(y) predicted_class = onecold(first(model(x, ps, st))) total_correct += sum(target_class .== predicted_class) total += length(target_class) end return total_correct / total end ``` ## Training {#Training} ```julia function train(model_function; cpu::Bool=false, kwargs...) dev = cpu ? cpu_device() : gpu_device() model, ps, st = create_model(model_function; dev, kwargs...) # Training train_dataloader, test_dataloader = dev(loadmnist(128, 0.9)) tstate = Training.TrainState(model, ps, st, Adam(0.001f0)) ### Lets train the model nepochs = 9 for epoch in 1:nepochs stime = time() for (x, y) in train_dataloader _, _, _, tstate = Training.single_train_step!( AutoZygote(), logitcrossentropy, (x, y), tstate ) end ttime = time() - stime tr_acc = accuracy(model, tstate.parameters, tstate.states, train_dataloader) * 100 te_acc = accuracy(model, tstate.parameters, tstate.states, test_dataloader) * 100 @printf "[%d/%d]\tTime %.4fs\tTraining Accuracy: %.5f%%\tTest \ Accuracy: %.5f%%\n" epoch nepochs ttime tr_acc te_acc end return nothing end train(NeuralODECompact) ``` ```julia train(NeuralODE) ``` We can also change the `sensealg` and train the model! `GaussAdjoint` allows you to use any arbitrary parameter structure and not just a flat vector (`ComponentArray`). ```julia train(NeuralODE; sensealg=GaussAdjoint(; autojacvec=ZygoteVJP()), use_named_tuple=true) ``` But remember some AD backends like `ReverseDiff` is not GPU compatible. For a model this size, you will notice that training time is significantly lower for training on CPU than on GPU. ```julia train(NeuralODE; sensealg=InterpolatingAdjoint(; autojacvec=ReverseDiffVJP()), cpu=true) ``` For completeness, let's also test out discrete sensitivities! ```julia train(NeuralODE; sensealg=ReverseDiffAdjoint(), cpu=true) ``` ## Alternate Implementation using Stateful Layer {#Alternate-Implementation-using-Stateful-Layer} Starting `v0.5.5`, Lux provides a `StatefulLuxLayer` which can be used to avoid the [`Box`ing of `st`](https://github.com/JuliaLang/julia/issues/15276). Using the `@compact` API avoids this problem entirely. ```julia struct StatefulNeuralODE{M<:Lux.AbstractLuxLayer,So,T,K} <: Lux.AbstractLuxWrapperLayer{:model} model::M solver::So tspan::T kwargs::K end function StatefulNeuralODE( model::Lux.AbstractLuxLayer; solver=Tsit5(), tspan=(0.0f0, 1.0f0), kwargs... ) return StatefulNeuralODE(model, solver, tspan, kwargs) end function (n::StatefulNeuralODE)(x, ps, st) st_model = StatefulLuxLayer(n.model, ps, st) dudt(u, p, t) = st_model(u, p) prob = ODEProblem{false}(ODEFunction{false}(dudt), x, n.tspan, ps) return solve(prob, n.solver; n.kwargs...), st_model.st end ``` ## Train the new Stateful Neural ODE {#Train-the-new-Stateful-Neural-ODE} ```julia train(StatefulNeuralODE) ``` We might not see a significant difference in the training time, but let us investigate the type stabilities of the layers. ## Type Stability {#Type-Stability} ```julia model, ps, st = create_model(NeuralODE) model_stateful, ps_stateful, st_stateful = create_model(StatefulNeuralODE) x = ones(Float32, 28, 28, 1, 3) |> gpu_device(); ``` NeuralODE is not type stable due to the boxing of `st` ```julia @code_warntype model(x, ps, st) ``` We avoid the problem entirely by using `StatefulNeuralODE` ```julia @code_warntype model_stateful(x, ps_stateful, st_stateful) ``` Note, that we still recommend using this layer internally and not exposing this as the default API to the users. Finally checking the compact model ```julia model_compact, ps_compact, st_compact = create_model(NeuralODECompact) @code_warntype model_compact(x, ps_compact, st_compact) ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/3_HyperNet.md --- # Training a HyperNetwork on MNIST and FashionMNIST {#Training-a-HyperNetwork-on-MNIST-and-FashionMNIST} ## Package Imports {#Package-Imports} ```julia using Lux, ComponentArrays, MLDatasets, MLUtils, OneHotArrays, Optimisers, Printf, Random, Reactant ``` ## Loading Datasets {#Loading-Datasets} ```julia function load_dataset( ::Type{dset}, n_train::Union{Nothing,Int}, n_eval::Union{Nothing,Int}, batchsize::Int ) where {dset} (; features, targets) = if n_train === nothing tmp = dset(:train) tmp[1:length(tmp)] else dset(:train)[1:n_train] end x_train, y_train = reshape(features, 28, 28, 1, :), onehotbatch(targets, 0:9) (; features, targets) = if n_eval === nothing tmp = dset(:test) tmp[1:length(tmp)] else dset(:test)[1:n_eval] end x_test, y_test = reshape(features, 28, 28, 1, :), onehotbatch(targets, 0:9) return ( DataLoader( (x_train, y_train); batchsize=min(batchsize, size(x_train, 4)), shuffle=true, partial=false, ), DataLoader( (x_test, y_test); batchsize=min(batchsize, size(x_test, 4)), shuffle=false, partial=false, ), ) end function load_datasets(batchsize=32) n_train = parse(Bool, get(ENV, "CI", "false")) ? 1024 : nothing n_eval = parse(Bool, get(ENV, "CI", "false")) ? 32 : nothing return load_dataset.((MNIST, FashionMNIST), n_train, n_eval, batchsize) end ``` ## Implement a HyperNet Layer {#Implement-a-HyperNet-Layer} ```julia function HyperNet(weight_generator::AbstractLuxLayer, core_network::AbstractLuxLayer) ca_axes = getaxes( ComponentArray(Lux.initialparameters(Random.default_rng(), core_network)) ) return @compact(; ca_axes, weight_generator, core_network, dispatch=:HyperNet) do (x, y) # Generate the weights ps_new = ComponentArray(vec(weight_generator(x)), ca_axes) @return core_network(y, ps_new) end end ``` Defining functions on the CompactLuxLayer requires some understanding of how the layer is structured, as such we don't recommend doing it unless you are familiar with the internals. In this case, we simply write it to ignore the initialization of the `core_network` parameters. ```julia function Lux.initialparameters(rng::AbstractRNG, hn::CompactLuxLayer{:HyperNet}) return (; weight_generator=Lux.initialparameters(rng, hn.layers.weight_generator)) end ``` ## Create and Initialize the HyperNet {#Create-and-Initialize-the-HyperNet} ```julia function create_model() core_network = Chain( Conv((3, 3), 1 => 16, relu; stride=2), Conv((3, 3), 16 => 32, relu; stride=2), Conv((3, 3), 32 => 64, relu; stride=2), GlobalMeanPool(), FlattenLayer(), Dense(64, 10), ) return HyperNet( Chain( Embedding(2 => 32), Dense(32, 64, relu), Dense(64, Lux.parameterlength(core_network)), ), core_network, ) end ``` ## Define Utility Functions {#Define-Utility-Functions} ```julia function accuracy(model, ps, st, dataloader, data_idx) total_correct, total = 0, 0 cdev = cpu_device() st = Lux.testmode(st) for (x, y) in dataloader ŷ, _ = model((data_idx, x), ps, st) target_class = y |> cdev |> onecold predicted_class = ŷ |> cdev |> onecold total_correct += sum(target_class .== predicted_class) total += length(target_class) end return total_correct / total end ``` ## Training {#Training} ```julia function train() dev = reactant_device(; force=true) model = create_model() dataloaders = load_datasets() |> dev Random.seed!(1234) ps, st = Lux.setup(Random.default_rng(), model) |> dev train_state = Training.TrainState(model, ps, st, Adam(0.0003f0)) x = first(first(dataloaders[1][1])) data_idx = ConcreteRNumber(1) model_compiled = @compile model((data_idx, x), ps, Lux.testmode(st)) ### Let's train the model nepochs = 50 for epoch in 1:nepochs, data_idx in 1:2 train_dataloader, test_dataloader = dev.(dataloaders[data_idx]) ### This allows us to trace the data index, else it will be embedded as a constant ### in the IR concrete_data_idx = ConcreteRNumber(data_idx) stime = time() for (x, y) in train_dataloader (_, _, _, train_state) = Training.single_train_step!( AutoEnzyme(), CrossEntropyLoss(; logits=Val(true)), ((concrete_data_idx, x), y), train_state; return_gradients=Val(false), ) end ttime = time() - stime train_acc = round( accuracy( model_compiled, train_state.parameters, train_state.states, train_dataloader, concrete_data_idx, ) * 100; digits=2, ) test_acc = round( accuracy( model_compiled, train_state.parameters, train_state.states, test_dataloader, concrete_data_idx, ) * 100; digits=2, ) data_name = data_idx == 1 ? "MNIST" : "FashionMNIST" @printf "[%3d/%3d]\t%12s\tTime %3.5fs\tTraining Accuracy: %3.2f%%\tTest \ Accuracy: %3.2f%%\n" epoch nepochs data_name ttime train_acc test_acc end println() test_acc_list = [0.0, 0.0] for data_idx in 1:2 train_dataloader, test_dataloader = dev.(dataloaders[data_idx]) concrete_data_idx = ConcreteRNumber(data_idx) train_acc = round( accuracy( model_compiled, train_state.parameters, train_state.states, train_dataloader, concrete_data_idx, ) * 100; digits=2, ) test_acc = round( accuracy( model_compiled, train_state.parameters, train_state.states, test_dataloader, concrete_data_idx, ) * 100; digits=2, ) data_name = data_idx == 1 ? "MNIST" : "FashionMNIST" @printf "[FINAL]\t%12s\tTraining Accuracy: %3.2f%%\tTest Accuracy: \ %3.2f%%\n" data_name train_acc test_acc test_acc_list[data_idx] = test_acc end return test_acc_list end test_acc_list = train() ``` ``` [ 1/ 50] MNIST Time 48.84021s Training Accuracy: 34.57% Test Accuracy: 37.50% [ 1/ 50] FashionMNIST Time 0.09202s Training Accuracy: 32.52% Test Accuracy: 43.75% [ 2/ 50] MNIST Time 0.09128s Training Accuracy: 36.33% Test Accuracy: 34.38% [ 2/ 50] FashionMNIST Time 0.08762s Training Accuracy: 46.19% Test Accuracy: 46.88% [ 3/ 50] MNIST Time 0.08880s Training Accuracy: 42.77% Test Accuracy: 28.12% [ 3/ 50] FashionMNIST Time 0.08956s Training Accuracy: 56.74% Test Accuracy: 56.25% [ 4/ 50] MNIST Time 0.08838s Training Accuracy: 51.17% Test Accuracy: 37.50% [ 4/ 50] FashionMNIST Time 0.09029s Training Accuracy: 63.48% Test Accuracy: 62.50% [ 5/ 50] MNIST Time 0.08769s Training Accuracy: 55.96% Test Accuracy: 43.75% [ 5/ 50] FashionMNIST Time 0.08950s Training Accuracy: 71.19% Test Accuracy: 53.12% [ 6/ 50] MNIST Time 0.08884s Training Accuracy: 63.09% Test Accuracy: 34.38% [ 6/ 50] FashionMNIST Time 0.08820s Training Accuracy: 74.32% Test Accuracy: 56.25% [ 7/ 50] MNIST Time 0.09790s Training Accuracy: 67.97% Test Accuracy: 50.00% [ 7/ 50] FashionMNIST Time 0.08662s Training Accuracy: 76.37% Test Accuracy: 62.50% [ 8/ 50] MNIST Time 0.08902s Training Accuracy: 74.80% Test Accuracy: 46.88% [ 8/ 50] FashionMNIST Time 0.10055s Training Accuracy: 81.93% Test Accuracy: 65.62% [ 9/ 50] MNIST Time 0.08639s Training Accuracy: 81.05% Test Accuracy: 50.00% [ 9/ 50] FashionMNIST Time 0.08597s Training Accuracy: 85.16% Test Accuracy: 62.50% [ 10/ 50] MNIST Time 0.09936s Training Accuracy: 82.42% Test Accuracy: 50.00% [ 10/ 50] FashionMNIST Time 0.08901s Training Accuracy: 87.11% Test Accuracy: 59.38% [ 11/ 50] MNIST Time 0.09011s Training Accuracy: 87.30% Test Accuracy: 50.00% [ 11/ 50] FashionMNIST Time 0.09869s Training Accuracy: 89.45% Test Accuracy: 68.75% [ 12/ 50] MNIST Time 0.09153s Training Accuracy: 89.75% Test Accuracy: 50.00% [ 12/ 50] FashionMNIST Time 0.08788s Training Accuracy: 90.82% Test Accuracy: 71.88% [ 13/ 50] MNIST Time 0.10019s Training Accuracy: 93.85% Test Accuracy: 59.38% [ 13/ 50] FashionMNIST Time 0.08920s Training Accuracy: 94.34% Test Accuracy: 68.75% [ 14/ 50] MNIST Time 0.09109s Training Accuracy: 93.95% Test Accuracy: 56.25% [ 14/ 50] FashionMNIST Time 0.09797s Training Accuracy: 95.51% Test Accuracy: 68.75% [ 15/ 50] MNIST Time 0.08859s Training Accuracy: 95.61% Test Accuracy: 62.50% [ 15/ 50] FashionMNIST Time 0.08875s Training Accuracy: 94.73% Test Accuracy: 71.88% [ 16/ 50] MNIST Time 0.08958s Training Accuracy: 97.95% Test Accuracy: 62.50% [ 16/ 50] FashionMNIST Time 0.08959s Training Accuracy: 96.48% Test Accuracy: 75.00% [ 17/ 50] MNIST Time 0.08911s Training Accuracy: 99.32% Test Accuracy: 62.50% [ 17/ 50] FashionMNIST Time 0.08885s Training Accuracy: 97.66% Test Accuracy: 71.88% [ 18/ 50] MNIST Time 0.08881s Training Accuracy: 99.51% Test Accuracy: 59.38% [ 18/ 50] FashionMNIST Time 0.08755s Training Accuracy: 97.46% Test Accuracy: 71.88% [ 19/ 50] MNIST Time 0.08810s Training Accuracy: 99.61% Test Accuracy: 59.38% [ 19/ 50] FashionMNIST Time 0.08844s Training Accuracy: 98.34% Test Accuracy: 68.75% [ 20/ 50] MNIST Time 0.08678s Training Accuracy: 99.90% Test Accuracy: 53.12% [ 20/ 50] FashionMNIST Time 0.10557s Training Accuracy: 98.54% Test Accuracy: 68.75% [ 21/ 50] MNIST Time 0.08847s Training Accuracy: 99.90% Test Accuracy: 62.50% [ 21/ 50] FashionMNIST Time 0.08642s Training Accuracy: 99.22% Test Accuracy: 75.00% [ 22/ 50] MNIST Time 0.08743s Training Accuracy: 99.90% Test Accuracy: 65.62% [ 22/ 50] FashionMNIST Time 0.08615s Training Accuracy: 99.32% Test Accuracy: 71.88% [ 23/ 50] MNIST Time 0.08697s Training Accuracy: 99.90% Test Accuracy: 65.62% [ 23/ 50] FashionMNIST Time 0.08838s Training Accuracy: 99.61% Test Accuracy: 71.88% [ 24/ 50] MNIST Time 0.08925s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 24/ 50] FashionMNIST Time 0.08695s Training Accuracy: 99.90% Test Accuracy: 71.88% [ 25/ 50] MNIST Time 0.08715s Training Accuracy: 100.00% Test Accuracy: 65.62% [ 25/ 50] FashionMNIST Time 0.09797s Training Accuracy: 99.90% Test Accuracy: 71.88% [ 26/ 50] MNIST Time 0.08676s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 26/ 50] FashionMNIST Time 0.08827s Training Accuracy: 100.00% Test Accuracy: 68.75% [ 27/ 50] MNIST Time 0.09931s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 27/ 50] FashionMNIST Time 0.08757s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 28/ 50] MNIST Time 0.08757s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 28/ 50] FashionMNIST Time 0.09637s Training Accuracy: 100.00% Test Accuracy: 78.12% [ 29/ 50] MNIST Time 0.08658s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 29/ 50] FashionMNIST Time 0.08639s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 30/ 50] MNIST Time 0.09771s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 30/ 50] FashionMNIST Time 0.08999s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 31/ 50] MNIST Time 0.09735s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 31/ 50] FashionMNIST Time 0.09602s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 32/ 50] MNIST Time 0.08743s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 32/ 50] FashionMNIST Time 0.08482s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 33/ 50] MNIST Time 0.08796s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 33/ 50] FashionMNIST Time 0.08779s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 34/ 50] MNIST Time 0.08921s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 34/ 50] FashionMNIST Time 0.08596s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 35/ 50] MNIST Time 0.08978s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 35/ 50] FashionMNIST Time 0.08639s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 36/ 50] MNIST Time 0.08804s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 36/ 50] FashionMNIST Time 0.09217s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 37/ 50] MNIST Time 0.08778s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 37/ 50] FashionMNIST Time 0.08569s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 38/ 50] MNIST Time 0.08629s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 38/ 50] FashionMNIST Time 0.08771s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 39/ 50] MNIST Time 0.08881s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 39/ 50] FashionMNIST Time 0.08830s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 40/ 50] MNIST Time 0.08885s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 40/ 50] FashionMNIST Time 0.08816s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 41/ 50] MNIST Time 0.09062s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 41/ 50] FashionMNIST Time 0.08828s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 42/ 50] MNIST Time 0.08959s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 42/ 50] FashionMNIST Time 0.09994s Training Accuracy: 100.00% Test Accuracy: 71.88% [ 43/ 50] MNIST Time 0.08845s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 43/ 50] FashionMNIST Time 0.08858s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 44/ 50] MNIST Time 0.10099s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 44/ 50] FashionMNIST Time 0.09119s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 45/ 50] MNIST Time 0.11106s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 45/ 50] FashionMNIST Time 0.10201s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 46/ 50] MNIST Time 0.08924s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 46/ 50] FashionMNIST Time 0.08997s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 47/ 50] MNIST Time 0.09913s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 47/ 50] FashionMNIST Time 0.08715s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 48/ 50] MNIST Time 0.08628s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 48/ 50] FashionMNIST Time 0.09982s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 49/ 50] MNIST Time 0.08876s Training Accuracy: 100.00% Test Accuracy: 62.50% [ 49/ 50] FashionMNIST Time 0.08811s Training Accuracy: 100.00% Test Accuracy: 75.00% [ 50/ 50] MNIST Time 0.08869s Training Accuracy: 100.00% Test Accuracy: 59.38% [ 50/ 50] FashionMNIST Time 0.08608s Training Accuracy: 100.00% Test Accuracy: 75.00% [FINAL] MNIST Training Accuracy: 100.00% Test Accuracy: 62.50% [FINAL] FashionMNIST Training Accuracy: 100.00% Test Accuracy: 75.00% ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/4_PINN2DPDE.md --- # Training a PINN on 2D PDE {#Training-a-PINN-on-2D-PDE} In this tutorial we will go over using a PINN to solve 2D PDEs. We will be using the system from [NeuralPDE Tutorials](https://docs.sciml.ai/NeuralPDE/stable/tutorials/gpu/). However, we will be using our custom loss function and use nested AD capabilities of Lux.jl. This is a demonstration of Lux.jl. For serious use cases of PINNs, please refer to the package: [NeuralPDE.jl](https://github.com/SciML/NeuralPDE.jl). ## Package Imports {#Package-Imports} ```julia using Lux, Optimisers, Random, Printf, Statistics, MLUtils, OnlineStats, CairoMakie, Reactant, Enzyme const xdev = reactant_device(; force=true) const cdev = cpu_device() ``` ## Problem Definition {#Problem-Definition} Since Lux supports efficient nested AD upto 2nd order, we will rewrite the problem with first order derivatives, so that we can compute the gradients of the loss using 2nd order AD. ## Define the Neural Networks {#Define-the-Neural-Networks} All the networks take 3 input variables and output a scalar value. Here, we will define a wrapper over the 3 networks, so that we can train them using [`Training.TrainState`](/api/Lux/utilities#Lux.Training.TrainState). ```julia struct PINN{M} <: AbstractLuxWrapperLayer{:model} model::M end function PINN(; hidden_dims::Int=32) return PINN( Chain( Dense(3 => hidden_dims, tanh), Dense(hidden_dims => hidden_dims, tanh), Dense(hidden_dims => hidden_dims, tanh), Dense(hidden_dims => 1), ), ) end ``` ## Define the Loss Functions {#Define-the-Loss-Functions} We will define a custom loss function to compute the loss using 2nd order AD. For that, first we'll need to define the derivatives of our model: ```julia function ∂u_∂t(model::StatefulLuxLayer, xyt::AbstractArray) return Enzyme.gradient(Enzyme.Reverse, sum ∘ model, xyt)[1][3, :] end function ∂u_∂x(model::StatefulLuxLayer, xyt::AbstractArray) return Enzyme.gradient(Enzyme.Reverse, sum ∘ model, xyt)[1][1, :] end function ∂u_∂y(model::StatefulLuxLayer, xyt::AbstractArray) return Enzyme.gradient(Enzyme.Reverse, sum ∘ model, xyt)[1][2, :] end function ∂²u_∂x²(model::StatefulLuxLayer, xyt::AbstractArray) return Enzyme.gradient(Enzyme.Reverse, sum ∘ ∂u_∂x, Enzyme.Const(model), xyt)[2][1, :] end function ∂²u_∂y²(model::StatefulLuxLayer, xyt::AbstractArray) return Enzyme.gradient(Enzyme.Reverse, sum ∘ ∂u_∂y, Enzyme.Const(model), xyt)[2][2, :] end ``` We will use the following loss function ```julia function physics_informed_loss_function(model::StatefulLuxLayer, xyt::AbstractArray) return mean(abs2, ∂u_∂t(model, xyt) .- ∂²u_∂x²(model, xyt) .- ∂²u_∂y²(model, xyt)) end ``` Additionally, we need to compute the loss with respect to the boundary conditions. ```julia function mse_loss_function( model::StatefulLuxLayer, target::AbstractArray, xyt::AbstractArray ) return MSELoss()(model(xyt), target) end function loss_function(model, ps, st, (xyt, target_data, xyt_bc, target_bc)) smodel = StatefulLuxLayer(model, ps, st) physics_loss = physics_informed_loss_function(smodel, xyt) data_loss = mse_loss_function(smodel, target_data, xyt) bc_loss = mse_loss_function(smodel, target_bc, xyt_bc) loss = physics_loss + data_loss + bc_loss return loss, smodel.st, (; physics_loss, data_loss, bc_loss) end ``` ## Generate the Data {#Generate-the-Data} We will generate some random data to train the model on. We will take data on a square spatial and temporal domain $x \in \[0, 2]$, $y \in \[0, 2]$, and $t \in \[0, 2]$. Typically, you want to be smarter about the sampling process, but for the sake of simplicity, we will skip that. ```julia analytical_solution(x, y, t) = @. exp(x + y) * cos(x + y + 4t) analytical_solution(xyt) = analytical_solution(xyt[1, :], xyt[2, :], xyt[3, :]) ``` ```julia grid_len = 16 grid = range(0.0f0, 2.0f0; length=grid_len) xyt = stack([[elem...] for elem in vec(collect(Iterators.product(grid, grid, grid)))]) target_data = reshape(analytical_solution(xyt), 1, :) bc_len = 512 x = collect(range(0.0f0, 2.0f0; length=bc_len)) y = collect(range(0.0f0, 2.0f0; length=bc_len)) t = collect(range(0.0f0, 2.0f0; length=bc_len)) xyt_bc = hcat( stack((x, y, zeros(Float32, bc_len)); dims=1), stack((zeros(Float32, bc_len), y, t); dims=1), stack((ones(Float32, bc_len) .* 2, y, t); dims=1), stack((x, zeros(Float32, bc_len), t); dims=1), stack((x, ones(Float32, bc_len) .* 2, t); dims=1), ) target_bc = reshape(analytical_solution(xyt_bc), 1, :) min_target_bc, max_target_bc = extrema(target_bc) min_data, max_data = extrema(target_data) min_pde_val, max_pde_val = min(min_data, min_target_bc), max(max_data, max_target_bc) xyt = (xyt .- minimum(xyt)) ./ (maximum(xyt) .- minimum(xyt)) xyt_bc = (xyt_bc .- minimum(xyt_bc)) ./ (maximum(xyt_bc) .- minimum(xyt_bc)) target_bc = (target_bc .- min_pde_val) ./ (max_pde_val - min_pde_val) target_data = (target_data .- min_pde_val) ./ (max_pde_val - min_pde_val) ``` ## Training {#Training} ```julia function train_model( xyt, target_data, xyt_bc, target_bc; seed::Int=0, maxiters::Int=50000, hidden_dims::Int=128, ) rng = Random.default_rng() Random.seed!(rng, seed) pinn = PINN(; hidden_dims) ps, st = Lux.setup(rng, pinn) |> xdev bc_dataloader = DataLoader((xyt_bc, target_bc); batchsize=128, shuffle=true, partial=false) |> xdev pde_dataloader = DataLoader((xyt, target_data); batchsize=128, shuffle=true, partial=false) |> xdev train_state = Training.TrainState(pinn, ps, st, Adam(0.005f0)) lr = i -> i < 5000 ? 0.005f0 : (i < 10000 ? 0.0005f0 : 0.00005f0) total_loss_tracker, physics_loss_tracker, data_loss_tracker, bc_loss_tracker = ntuple( _ -> OnlineStats.CircBuff(Float32, 32; rev=true), 4 ) iter = 1 for ((xyt_batch, target_data_batch), (xyt_bc_batch, target_bc_batch)) in zip(Iterators.cycle(pde_dataloader), Iterators.cycle(bc_dataloader)) Optimisers.adjust!(train_state, lr(iter)) _, loss, stats, train_state = Training.single_train_step!( AutoEnzyme(), loss_function, (xyt_batch, target_data_batch, xyt_bc_batch, target_bc_batch), train_state; return_gradients=Val(false), ) fit!(total_loss_tracker, Float32(loss)) fit!(physics_loss_tracker, Float32(stats.physics_loss)) fit!(data_loss_tracker, Float32(stats.data_loss)) fit!(bc_loss_tracker, Float32(stats.bc_loss)) mean_loss = mean(OnlineStats.value(total_loss_tracker)) mean_physics_loss = mean(OnlineStats.value(physics_loss_tracker)) mean_data_loss = mean(OnlineStats.value(data_loss_tracker)) mean_bc_loss = mean(OnlineStats.value(bc_loss_tracker)) isnan(loss) && throw(ArgumentError("NaN Loss Detected")) if iter % 1000 == 1 || iter == maxiters @printf( "Iteration: [%6d/%6d] \t Loss: %.9f (%.9f) \t Physics Loss: %.9f \ (%.9f) \t Data Loss: %.9f (%.9f) \t BC \ Loss: %.9f (%.9f)\n", iter, maxiters, loss, mean_loss, stats.physics_loss, mean_physics_loss, stats.data_loss, mean_data_loss, stats.bc_loss, mean_bc_loss ) end iter += 1 iter ≥ maxiters && break end return StatefulLuxLayer(pinn, cdev(train_state.parameters), cdev(train_state.states)) end trained_model = train_model(xyt, target_data, xyt_bc, target_bc) ``` ``` Iteration: [ 1/ 50000] Loss: 20.523933411 (20.523933411) Physics Loss: 16.931318283 (16.931318283) Data Loss: 2.007483006 (2.007483006) BC Loss: 1.585133076 (1.585133076) Iteration: [ 1001/ 50000] Loss: 0.017368633 (0.019241147) Physics Loss: 0.000384352 (0.000523631) Data Loss: 0.005318501 (0.007538573) BC Loss: 0.011665781 (0.011178940) Iteration: [ 2001/ 50000] Loss: 0.015431664 (0.018665710) Physics Loss: 0.001248588 (0.001662086) Data Loss: 0.004322519 (0.006408236) BC Loss: 0.009860558 (0.010595387) Iteration: [ 3001/ 50000] Loss: 0.015749730 (0.015216020) Physics Loss: 0.000569911 (0.001279066) Data Loss: 0.004014889 (0.004232434) BC Loss: 0.011164930 (0.009704520) Iteration: [ 4001/ 50000] Loss: 0.009721411 (0.008720266) Physics Loss: 0.002389135 (0.003386307) Data Loss: 0.003174899 (0.002104989) BC Loss: 0.004157377 (0.003228969) Iteration: [ 5001/ 50000] Loss: 0.004499406 (0.004708527) Physics Loss: 0.001959276 (0.001955591) Data Loss: 0.001685009 (0.001497321) BC Loss: 0.000855121 (0.001255615) Iteration: [ 6001/ 50000] Loss: 0.001065215 (0.001272577) Physics Loss: 0.000281185 (0.000287863) Data Loss: 0.000585776 (0.000744654) BC Loss: 0.000198253 (0.000240060) Iteration: [ 7001/ 50000] Loss: 0.001347864 (0.000929980) Physics Loss: 0.000297512 (0.000298550) Data Loss: 0.000953921 (0.000510000) BC Loss: 0.000096431 (0.000121430) Iteration: [ 8001/ 50000] Loss: 0.001117045 (0.000741427) Physics Loss: 0.000746706 (0.000274790) Data Loss: 0.000305004 (0.000386474) BC Loss: 0.000065335 (0.000080163) Iteration: [ 9001/ 50000] Loss: 0.001055349 (0.001118247) Physics Loss: 0.000357186 (0.000671388) Data Loss: 0.000604138 (0.000345194) BC Loss: 0.000094025 (0.000101665) Iteration: [ 10001/ 50000] Loss: 0.000577895 (0.000586777) Physics Loss: 0.000227950 (0.000239104) Data Loss: 0.000296419 (0.000296933) BC Loss: 0.000053526 (0.000050740) Iteration: [ 11001/ 50000] Loss: 0.000387099 (0.000371386) Physics Loss: 0.000160282 (0.000064884) Data Loss: 0.000181874 (0.000270512) BC Loss: 0.000044943 (0.000035990) Iteration: [ 12001/ 50000] Loss: 0.000264985 (0.000347079) Physics Loss: 0.000050845 (0.000063801) Data Loss: 0.000172765 (0.000247937) BC Loss: 0.000041375 (0.000035341) Iteration: [ 13001/ 50000] Loss: 0.000303605 (0.000328298) Physics Loss: 0.000061343 (0.000068644) Data Loss: 0.000208084 (0.000226532) BC Loss: 0.000034178 (0.000033121) Iteration: [ 14001/ 50000] Loss: 0.000384829 (0.000331977) Physics Loss: 0.000068277 (0.000069249) Data Loss: 0.000282153 (0.000234252) BC Loss: 0.000034399 (0.000028476) Iteration: [ 15001/ 50000] Loss: 0.000239852 (0.000288137) Physics Loss: 0.000043116 (0.000057610) Data Loss: 0.000166922 (0.000199768) BC Loss: 0.000029814 (0.000030758) Iteration: [ 16001/ 50000] Loss: 0.000223553 (0.000286669) Physics Loss: 0.000050308 (0.000056936) Data Loss: 0.000142008 (0.000201499) BC Loss: 0.000031237 (0.000028234) Iteration: [ 17001/ 50000] Loss: 0.000421828 (0.000296886) Physics Loss: 0.000095925 (0.000067681) Data Loss: 0.000300998 (0.000201824) BC Loss: 0.000024905 (0.000027381) Iteration: [ 18001/ 50000] Loss: 0.000211644 (0.000281161) Physics Loss: 0.000038092 (0.000057524) Data Loss: 0.000140487 (0.000195981) BC Loss: 0.000033065 (0.000027656) Iteration: [ 19001/ 50000] Loss: 0.000211127 (0.000277151) Physics Loss: 0.000053733 (0.000055449) Data Loss: 0.000138345 (0.000197246) BC Loss: 0.000019049 (0.000024456) Iteration: [ 20001/ 50000] Loss: 0.000310542 (0.000249381) Physics Loss: 0.000061448 (0.000046612) Data Loss: 0.000229472 (0.000180468) BC Loss: 0.000019622 (0.000022300) Iteration: [ 21001/ 50000] Loss: 0.000297450 (0.000248106) Physics Loss: 0.000059162 (0.000053576) Data Loss: 0.000216021 (0.000171196) BC Loss: 0.000022266 (0.000023334) Iteration: [ 22001/ 50000] Loss: 0.000164182 (0.000235557) Physics Loss: 0.000028181 (0.000045489) Data Loss: 0.000109619 (0.000167632) BC Loss: 0.000026382 (0.000022436) Iteration: [ 23001/ 50000] Loss: 0.000232555 (0.000246731) Physics Loss: 0.000037212 (0.000047894) Data Loss: 0.000171760 (0.000177439) BC Loss: 0.000023583 (0.000021399) Iteration: [ 24001/ 50000] Loss: 0.000292702 (0.000249951) Physics Loss: 0.000055334 (0.000055020) Data Loss: 0.000219019 (0.000170300) BC Loss: 0.000018350 (0.000024631) Iteration: [ 25001/ 50000] Loss: 0.000206420 (0.000227030) Physics Loss: 0.000044298 (0.000038203) Data Loss: 0.000143160 (0.000167706) BC Loss: 0.000018962 (0.000021121) Iteration: [ 26001/ 50000] Loss: 0.000229478 (0.000242765) Physics Loss: 0.000053504 (0.000056142) Data Loss: 0.000154916 (0.000164058) BC Loss: 0.000021058 (0.000022566) Iteration: [ 27001/ 50000] Loss: 0.000224114 (0.000242052) Physics Loss: 0.000045716 (0.000054974) Data Loss: 0.000152396 (0.000164748) BC Loss: 0.000026002 (0.000022330) Iteration: [ 28001/ 50000] Loss: 0.000222216 (0.000215397) Physics Loss: 0.000050252 (0.000036081) Data Loss: 0.000154347 (0.000157962) BC Loss: 0.000017617 (0.000021354) Iteration: [ 29001/ 50000] Loss: 0.000210704 (0.000233809) Physics Loss: 0.000035138 (0.000053348) Data Loss: 0.000143131 (0.000158839) BC Loss: 0.000032434 (0.000021622) Iteration: [ 30001/ 50000] Loss: 0.000207375 (0.000226769) Physics Loss: 0.000026813 (0.000043914) Data Loss: 0.000155654 (0.000161216) BC Loss: 0.000024908 (0.000021639) Iteration: [ 31001/ 50000] Loss: 0.000269214 (0.000219518) Physics Loss: 0.000039355 (0.000040133) Data Loss: 0.000209013 (0.000158317) BC Loss: 0.000020846 (0.000021067) Iteration: [ 32001/ 50000] Loss: 0.000211419 (0.000210957) Physics Loss: 0.000045159 (0.000038796) Data Loss: 0.000146814 (0.000151931) BC Loss: 0.000019446 (0.000020230) Iteration: [ 33001/ 50000] Loss: 0.000181482 (0.000202039) Physics Loss: 0.000025003 (0.000032754) Data Loss: 0.000134261 (0.000148478) BC Loss: 0.000022218 (0.000020807) Iteration: [ 34001/ 50000] Loss: 0.000186912 (0.000197588) Physics Loss: 0.000025085 (0.000032541) Data Loss: 0.000141114 (0.000145119) BC Loss: 0.000020713 (0.000019928) Iteration: [ 35001/ 50000] Loss: 0.000136710 (0.000217637) Physics Loss: 0.000018034 (0.000050786) Data Loss: 0.000097950 (0.000147103) BC Loss: 0.000020726 (0.000019747) Iteration: [ 36001/ 50000] Loss: 0.000158078 (0.000199009) Physics Loss: 0.000024817 (0.000033946) Data Loss: 0.000114559 (0.000146315) BC Loss: 0.000018701 (0.000018748) Iteration: [ 37001/ 50000] Loss: 0.000302340 (0.000190165) Physics Loss: 0.000066689 (0.000025736) Data Loss: 0.000217540 (0.000145103) BC Loss: 0.000018111 (0.000019326) Iteration: [ 38001/ 50000] Loss: 0.000245146 (0.000204957) Physics Loss: 0.000030379 (0.000036335) Data Loss: 0.000189912 (0.000150965) BC Loss: 0.000024854 (0.000017657) Iteration: [ 39001/ 50000] Loss: 0.000165623 (0.000197576) Physics Loss: 0.000026971 (0.000032796) Data Loss: 0.000122209 (0.000145698) BC Loss: 0.000016443 (0.000019082) Iteration: [ 40001/ 50000] Loss: 0.000171607 (0.000195268) Physics Loss: 0.000026284 (0.000031357) Data Loss: 0.000123539 (0.000143961) BC Loss: 0.000021784 (0.000019949) Iteration: [ 41001/ 50000] Loss: 0.000156535 (0.000193990) Physics Loss: 0.000019421 (0.000028832) Data Loss: 0.000119100 (0.000146621) BC Loss: 0.000018014 (0.000018538) Iteration: [ 42001/ 50000] Loss: 0.000175756 (0.000190178) Physics Loss: 0.000023863 (0.000026718) Data Loss: 0.000135717 (0.000144790) BC Loss: 0.000016176 (0.000018671) Iteration: [ 43001/ 50000] Loss: 0.000191984 (0.000191988) Physics Loss: 0.000026218 (0.000027192) Data Loss: 0.000150546 (0.000144479) BC Loss: 0.000015221 (0.000020318) Iteration: [ 44001/ 50000] Loss: 0.000171279 (0.000191033) Physics Loss: 0.000011855 (0.000027447) Data Loss: 0.000144759 (0.000142707) BC Loss: 0.000014665 (0.000020879) Iteration: [ 45001/ 50000] Loss: 0.000266982 (0.000207141) Physics Loss: 0.000038990 (0.000034863) Data Loss: 0.000201324 (0.000150720) BC Loss: 0.000026669 (0.000021558) Iteration: [ 46001/ 50000] Loss: 0.000201456 (0.000188971) Physics Loss: 0.000026719 (0.000033167) Data Loss: 0.000155729 (0.000138071) BC Loss: 0.000019007 (0.000017732) Iteration: [ 47001/ 50000] Loss: 0.000180822 (0.000185797) Physics Loss: 0.000026415 (0.000028996) Data Loss: 0.000137763 (0.000138725) BC Loss: 0.000016644 (0.000018075) Iteration: [ 48001/ 50000] Loss: 0.000188553 (0.000178357) Physics Loss: 0.000028802 (0.000022570) Data Loss: 0.000141659 (0.000137397) BC Loss: 0.000018092 (0.000018390) Iteration: [ 49001/ 50000] Loss: 0.000172877 (0.000196456) Physics Loss: 0.000030594 (0.000037507) Data Loss: 0.000120721 (0.000140691) BC Loss: 0.000021562 (0.000018257) ``` ## Visualizing the Results {#Visualizing-the-Results} ```julia ts, xs, ys = 0.0f0:0.05f0:2.0f0, 0.0f0:0.02f0:2.0f0, 0.0f0:0.02f0:2.0f0 grid = stack([[elem...] for elem in vec(collect(Iterators.product(xs, ys, ts)))]) u_real = reshape(analytical_solution(grid), length(xs), length(ys), length(ts)) grid_normalized = (grid .- minimum(grid)) ./ (maximum(grid) .- minimum(grid)) u_pred = reshape(trained_model(grid_normalized), length(xs), length(ys), length(ts)) u_pred = u_pred .* (max_pde_val - min_pde_val) .+ min_pde_val begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y") errs = [abs.(u_pred[:, :, i] .- u_real[:, :, i]) for i in 1:length(ts)] Colorbar(fig[1, 2]; limits=extrema(stack(errs))) CairoMakie.record(fig, "pinn_nested_ad.gif", 1:length(ts); framerate=10) do i ax.title = "Abs. Predictor Error | Time: $(ts[i])" err = errs[i] contour!(ax, xs, ys, err; levels=10, linewidth=2) heatmap!(ax, xs, ys, err) return fig end fig end ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, icelake-server) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/5_ConvolutionalVAE.md --- # Convolutional VAE for MNIST {#Convolutional-VAE-Tutorial} Convolutional variational autoencoder (CVAE) implementation in MLX using MNIST. This is based on the [CVAE implementation in MLX](https://github.com/ml-explore/mlx-examples/blob/main/cvae/). ```julia using Lux, Reactant, MLDatasets, Random, Statistics, Enzyme, MLUtils, DataAugmentation, ConcreteStructs, OneHotArrays, ImageShow, Images, Printf, Optimisers const xdev = reactant_device(; force=true) const cdev = cpu_device() const IN_VSCODE = isdefined(Main, :VSCodeServer) ``` ``` false ``` ## Model Definition {#Model-Definition} First we will define the encoder.It maps the input to a normal distribution in latent space and sample a latent vector from that distribution. ```julia function cvae_encoder( rng=Random.default_rng(); num_latent_dims::Int, image_shape::Dims{3}, max_num_filters::Int, ) flattened_dim = prod(image_shape[1:2] .÷ 8) * max_num_filters return @compact(; embed=Chain( Chain( Conv((3, 3), image_shape[3] => max_num_filters ÷ 4; stride=2, pad=1), BatchNorm(max_num_filters ÷ 4, leakyrelu), ), Chain( Conv((3, 3), max_num_filters ÷ 4 => max_num_filters ÷ 2; stride=2, pad=1), BatchNorm(max_num_filters ÷ 2, leakyrelu), ), Chain( Conv((3, 3), max_num_filters ÷ 2 => max_num_filters; stride=2, pad=1), BatchNorm(max_num_filters, leakyrelu), ), FlattenLayer(), ), proj_mu=Dense(flattened_dim, num_latent_dims; init_bias=zeros32), proj_log_var=Dense(flattened_dim, num_latent_dims; init_bias=zeros32), rng ) do x y = embed(x) μ = proj_mu(y) logσ² = proj_log_var(y) T = eltype(logσ²) logσ² = clamp.(logσ², -T(20.0f0), T(10.0f0)) σ = exp.(logσ² .* T(0.5)) # Generate a tensor of random values from a normal distribution ϵ = randn_like(Lux.replicate(rng), σ) # Reparameterization trick to backpropagate through sampling z = ϵ .* σ .+ μ @return z, μ, logσ² end end ``` Similarly we define the decoder. ```julia function cvae_decoder(; num_latent_dims::Int, image_shape::Dims{3}, max_num_filters::Int) flattened_dim = prod(image_shape[1:2] .÷ 8) * max_num_filters return @compact(; linear=Dense(num_latent_dims, flattened_dim), upchain=Chain( Chain( Upsample(2), Conv((3, 3), max_num_filters => max_num_filters ÷ 2; stride=1, pad=1), BatchNorm(max_num_filters ÷ 2, leakyrelu), ), Chain( Upsample(2), Conv((3, 3), max_num_filters ÷ 2 => max_num_filters ÷ 4; stride=1, pad=1), BatchNorm(max_num_filters ÷ 4, leakyrelu), ), Chain( Upsample(2), Conv( (3, 3), max_num_filters ÷ 4 => image_shape[3], sigmoid; stride=1, pad=1 ), ), ), max_num_filters ) do x y = linear(x) img = reshape(y, image_shape[1] ÷ 8, image_shape[2] ÷ 8, max_num_filters, :) @return upchain(img) end end @concrete struct CVAE <: AbstractLuxContainerLayer{(:encoder, :decoder)} encoder <: AbstractLuxLayer decoder <: AbstractLuxLayer end function CVAE( rng=Random.default_rng(); num_latent_dims::Int, image_shape::Dims{3}, max_num_filters::Int, ) decoder = cvae_decoder(; num_latent_dims, image_shape, max_num_filters) encoder = cvae_encoder(rng; num_latent_dims, image_shape, max_num_filters) return CVAE(encoder, decoder) end function (cvae::CVAE)(x, ps, st) (z, μ, logσ²), st_enc = cvae.encoder(x, ps.encoder, st.encoder) x_rec, st_dec = cvae.decoder(z, ps.decoder, st.decoder) return (x_rec, μ, logσ²), (; encoder=st_enc, decoder=st_dec) end function encode(cvae::CVAE, x, ps, st) (z, _, _), st_enc = cvae.encoder(x, ps.encoder, st.encoder) return z, (; encoder=st_enc, st.decoder) end function decode(cvae::CVAE, z, ps, st) x_rec, st_dec = cvae.decoder(z, ps.decoder, st.decoder) return x_rec, (; decoder=st_dec, st.encoder) end ``` ## Loading MNIST {#Loading-MNIST} ```julia @concrete struct TensorDataset dataset transform total_samples::Int end Base.length(ds::TensorDataset) = ds.total_samples function Base.getindex(ds::TensorDataset, idxs::Union{Vector{<:Integer},AbstractRange}) img = Image.(eachslice(convert2image(ds.dataset, idxs); dims=3)) return stack(parent ∘ itemdata ∘ Base.Fix1(apply, ds.transform), img) end function loadmnist(batchsize, image_size::Dims{2}) # Load MNIST: Only 1500 for demonstration purposes on CI train_dataset = MNIST(; split=:train) N = parse(Bool, get(ENV, "CI", "false")) ? 5000 : length(train_dataset) train_transform = ScaleKeepAspect(image_size) |> ImageToTensor() trainset = TensorDataset(train_dataset, train_transform, N) trainloader = DataLoader(trainset; batchsize, shuffle=true, partial=false) return trainloader end ``` ## Helper Functions {#Helper-Functions} Generate an Image Grid from a list of images ```julia function create_image_grid(imgs::AbstractArray, grid_rows::Int, grid_cols::Int) total_images = grid_rows * grid_cols imgs = map(eachslice(imgs[:, :, :, 1:total_images]; dims=4)) do img cimg = if size(img, 3) == 1 colorview(Gray, view(img, :, :, 1)) else colorview(RGB, permutedims(img, (3, 1, 2))) end return cimg' end return create_image_grid(imgs, grid_rows, grid_cols) end function create_image_grid(images::Vector, grid_rows::Int, grid_cols::Int) # Check if the number of images matches the grid total_images = grid_rows * grid_cols @assert length(images) == total_images # Get the size of a single image (assuming all images are the same size) img_height, img_width = size(images[1]) # Create a blank grid canvas grid_height = img_height * grid_rows grid_width = img_width * grid_cols grid_canvas = similar(images[1], grid_height, grid_width) # Place each image in the correct position on the canvas for idx in 1:total_images row = div(idx - 1, grid_cols) + 1 col = mod(idx - 1, grid_cols) + 1 start_row = (row - 1) * img_height + 1 start_col = (col - 1) * img_width + 1 grid_canvas[start_row:(start_row + img_height - 1), start_col:(start_col + img_width - 1)] .= images[idx] end return grid_canvas end function loss_function(model, ps, st, X) (y, μ, logσ²), st = model(X, ps, st) reconstruction_loss = MSELoss(; agg=sum)(y, X) kldiv_loss = -sum(1 .+ logσ² .- μ .^ 2 .- exp.(logσ²)) / 2 loss = reconstruction_loss + kldiv_loss return loss, st, (; y, μ, logσ², reconstruction_loss, kldiv_loss) end function generate_images( model, ps, st; num_samples::Int=128, num_latent_dims::Int, decode_compiled=nothing ) z = get_device((ps, st))(randn(Float32, num_latent_dims, num_samples)) if decode_compiled === nothing images, _ = decode(model, z, ps, Lux.testmode(st)) else images, _ = decode_compiled(model, z, ps, Lux.testmode(st)) images = cpu_device()(images) end return create_image_grid(images, 8, num_samples ÷ 8) end function reconstruct_images(model, ps, st, X) (recon, _, _), _ = model(X, ps, Lux.testmode(st)) recon = cpu_device()(recon) return create_image_grid(recon, 8, size(X, ndims(X)) ÷ 8) end ``` ``` reconstruct_images (generic function with 1 method) ``` ## Training the Model {#Training-the-Model} ```julia function main(; batchsize=128, image_size=(64, 64), num_latent_dims=8, max_num_filters=64, seed=0, epochs=50, weight_decay=1.0e-5, learning_rate=1.0e-3, num_samples=batchsize, ) rng = Xoshiro() Random.seed!(rng, seed) cvae = CVAE(rng; num_latent_dims, image_shape=(image_size..., 1), max_num_filters) ps, st = Lux.setup(rng, cvae) |> xdev z = xdev(randn(Float32, num_latent_dims, num_samples)) decode_compiled = @compile decode(cvae, z, ps, Lux.testmode(st)) x = randn(Float32, image_size..., 1, batchsize) |> xdev cvae_compiled = @compile cvae(x, ps, Lux.testmode(st)) train_dataloader = loadmnist(batchsize, image_size) |> xdev opt = AdamW(; eta=learning_rate, lambda=weight_decay) train_state = Training.TrainState(cvae, ps, st, opt) @printf "Total Trainable Parameters: %0.4f M\n" (Lux.parameterlength(ps) / 1.0e6) empty_row, model_img_full = nothing, nothing for epoch in 1:epochs loss_total = 0.0f0 total_samples = 0 start_time = time() for (i, X) in enumerate(train_dataloader) (_, loss, _, train_state) = Training.single_train_step!( AutoEnzyme(), loss_function, X, train_state; return_gradients=Val(false) ) loss_total += loss total_samples += size(X, ndims(X)) if i % 250 == 0 || i == length(train_dataloader) throughput = total_samples / (time() - start_time) @printf "Epoch %d, Iter %d, Loss: %.7f, Throughput: %.6f im/s\n" epoch i loss throughput end end total_time = time() - start_time train_loss = loss_total / length(train_dataloader) throughput = total_samples / total_time @printf "Epoch %d, Train Loss: %.7f, Time: %.4fs, Throughput: %.6f im/s\n" epoch train_loss total_time throughput if IN_VSCODE || epoch == epochs recon_images = reconstruct_images( cvae_compiled, train_state.parameters, train_state.states, first(train_dataloader), ) gen_images = generate_images( cvae, train_state.parameters, train_state.states; num_samples, num_latent_dims, decode_compiled, ) if empty_row === nothing empty_row = similar(gen_images, image_size[1], size(gen_images, 2)) fill!(empty_row, 0) end model_img_full = vcat(recon_images, empty_row, gen_images) IN_VSCODE && display(model_img_full) end end return model_img_full end img = main() ``` ``` Total Trainable Parameters: 0.1493 M Epoch 1, Iter 39, Loss: 23669.0214844, Throughput: 10.124549 im/s Epoch 1, Train Loss: 39975.1601562, Time: 493.3634s, Throughput: 10.118301 im/s Epoch 2, Iter 39, Loss: 18302.7792969, Throughput: 68.049236 im/s Epoch 2, Train Loss: 20375.8925781, Time: 73.3589s, Throughput: 68.049023 im/s Epoch 3, Iter 39, Loss: 16283.4511719, Throughput: 68.054038 im/s Epoch 3, Train Loss: 16709.9121094, Time: 73.3537s, Throughput: 68.053838 im/s Epoch 4, Iter 39, Loss: 13466.8798828, Throughput: 68.032190 im/s Epoch 4, Train Loss: 15122.2822266, Time: 73.3772s, Throughput: 68.032016 im/s Epoch 5, Iter 39, Loss: 14304.7539062, Throughput: 68.043028 im/s Epoch 5, Train Loss: 14137.8144531, Time: 73.3655s, Throughput: 68.042864 im/s Epoch 6, Iter 39, Loss: 12600.3085938, Throughput: 68.001696 im/s Epoch 6, Train Loss: 13351.3906250, Time: 73.4102s, Throughput: 68.001472 im/s Epoch 7, Iter 39, Loss: 12405.7949219, Throughput: 69.244548 im/s Epoch 7, Train Loss: 12894.4794922, Time: 72.0925s, Throughput: 69.244354 im/s Epoch 8, Iter 39, Loss: 12413.3720703, Throughput: 68.602105 im/s Epoch 8, Train Loss: 12526.3232422, Time: 72.7676s, Throughput: 68.601936 im/s Epoch 9, Iter 39, Loss: 11890.4316406, Throughput: 68.033301 im/s Epoch 9, Train Loss: 12207.2285156, Time: 73.3760s, Throughput: 68.033115 im/s Epoch 10, Iter 39, Loss: 12109.4013672, Throughput: 67.690280 im/s Epoch 10, Train Loss: 11937.8769531, Time: 73.7479s, Throughput: 67.690100 im/s Epoch 11, Iter 39, Loss: 11775.3535156, Throughput: 68.061594 im/s Epoch 11, Train Loss: 11712.8369141, Time: 73.3455s, Throughput: 68.061436 im/s Epoch 12, Iter 39, Loss: 11403.8300781, Throughput: 68.355921 im/s Epoch 12, Train Loss: 11601.0283203, Time: 73.0297s, Throughput: 68.355734 im/s Epoch 13, Iter 39, Loss: 11250.1308594, Throughput: 68.534802 im/s Epoch 13, Train Loss: 11406.9287109, Time: 72.8391s, Throughput: 68.534622 im/s Epoch 14, Iter 39, Loss: 11398.3183594, Throughput: 68.165342 im/s Epoch 14, Train Loss: 11303.6406250, Time: 73.2339s, Throughput: 68.165161 im/s Epoch 15, Iter 39, Loss: 11767.1406250, Throughput: 68.562296 im/s Epoch 15, Train Loss: 11192.1826172, Time: 72.8099s, Throughput: 68.562135 im/s Epoch 16, Iter 39, Loss: 10836.2226562, Throughput: 69.231020 im/s Epoch 16, Train Loss: 10999.6738281, Time: 72.1066s, Throughput: 69.230829 im/s Epoch 17, Iter 39, Loss: 10251.7041016, Throughput: 69.279984 im/s Epoch 17, Train Loss: 10919.7597656, Time: 72.0556s, Throughput: 69.279818 im/s Epoch 18, Iter 39, Loss: 10644.7070312, Throughput: 69.279546 im/s Epoch 18, Train Loss: 10823.3603516, Time: 72.0561s, Throughput: 69.279341 im/s Epoch 19, Iter 39, Loss: 11289.3828125, Throughput: 68.848847 im/s Epoch 19, Train Loss: 10718.5107422, Time: 72.5068s, Throughput: 68.848679 im/s Epoch 20, Iter 39, Loss: 10420.6542969, Throughput: 68.845076 im/s Epoch 20, Train Loss: 10539.7568359, Time: 72.5108s, Throughput: 68.844940 im/s Epoch 21, Iter 39, Loss: 10542.1757812, Throughput: 68.777805 im/s Epoch 21, Train Loss: 10445.1005859, Time: 72.5817s, Throughput: 68.777634 im/s Epoch 22, Iter 39, Loss: 10361.1679688, Throughput: 68.743540 im/s Epoch 22, Train Loss: 10374.2812500, Time: 72.6179s, Throughput: 68.743382 im/s Epoch 23, Iter 39, Loss: 10557.1943359, Throughput: 68.776692 im/s Epoch 23, Train Loss: 10406.1728516, Time: 72.5829s, Throughput: 68.776532 im/s Epoch 24, Iter 39, Loss: 10106.6113281, Throughput: 68.653364 im/s Epoch 24, Train Loss: 10377.8544922, Time: 72.7133s, Throughput: 68.653187 im/s Epoch 25, Iter 39, Loss: 10092.6054688, Throughput: 68.915060 im/s Epoch 25, Train Loss: 10258.3574219, Time: 72.4372s, Throughput: 68.914895 im/s Epoch 26, Iter 39, Loss: 10645.6298828, Throughput: 68.927061 im/s Epoch 26, Train Loss: 10262.6572266, Time: 72.4246s, Throughput: 68.926895 im/s Epoch 27, Iter 39, Loss: 10337.3339844, Throughput: 68.412689 im/s Epoch 27, Train Loss: 10140.3837891, Time: 72.9691s, Throughput: 68.412527 im/s Epoch 28, Iter 39, Loss: 10293.0546875, Throughput: 68.378347 im/s Epoch 28, Train Loss: 10080.1660156, Time: 73.0057s, Throughput: 68.378184 im/s Epoch 29, Iter 39, Loss: 10612.8369141, Throughput: 68.715038 im/s Epoch 29, Train Loss: 10073.8808594, Time: 72.6480s, Throughput: 68.714876 im/s Epoch 30, Iter 39, Loss: 10239.8535156, Throughput: 68.469637 im/s Epoch 30, Train Loss: 9966.2177734, Time: 72.9084s, Throughput: 68.469470 im/s Epoch 31, Iter 39, Loss: 10710.3222656, Throughput: 68.559571 im/s Epoch 31, Train Loss: 10012.6406250, Time: 72.8128s, Throughput: 68.559358 im/s Epoch 32, Iter 39, Loss: 9745.9951172, Throughput: 68.676341 im/s Epoch 32, Train Loss: 9882.7548828, Time: 72.6890s, Throughput: 68.676143 im/s Epoch 33, Iter 39, Loss: 9787.8847656, Throughput: 69.333868 im/s Epoch 33, Train Loss: 9899.9130859, Time: 71.9996s, Throughput: 69.333696 im/s Epoch 34, Iter 39, Loss: 10539.0322266, Throughput: 68.515452 im/s Epoch 34, Train Loss: 9837.1962891, Time: 72.8597s, Throughput: 68.515289 im/s Epoch 35, Iter 39, Loss: 9744.6542969, Throughput: 67.934433 im/s Epoch 35, Train Loss: 9805.7861328, Time: 73.4828s, Throughput: 67.934235 im/s Epoch 36, Iter 39, Loss: 9546.8281250, Throughput: 67.951674 im/s Epoch 36, Train Loss: 9718.6044922, Time: 73.4642s, Throughput: 67.951502 im/s Epoch 37, Iter 39, Loss: 10119.2109375, Throughput: 68.748607 im/s Epoch 37, Train Loss: 9725.8994141, Time: 72.6126s, Throughput: 68.748417 im/s Epoch 38, Iter 39, Loss: 9806.8525391, Throughput: 67.718084 im/s Epoch 38, Train Loss: 9681.8994141, Time: 73.7175s, Throughput: 67.717933 im/s Epoch 39, Iter 39, Loss: 9570.3046875, Throughput: 67.956688 im/s Epoch 39, Train Loss: 9612.1855469, Time: 73.4587s, Throughput: 67.956508 im/s Epoch 40, Iter 39, Loss: 9559.7421875, Throughput: 68.146191 im/s Epoch 40, Train Loss: 9571.6318359, Time: 73.2545s, Throughput: 68.146002 im/s Epoch 41, Iter 39, Loss: 9602.0146484, Throughput: 68.271227 im/s Epoch 41, Train Loss: 9613.8828125, Time: 73.1203s, Throughput: 68.271059 im/s Epoch 42, Iter 39, Loss: 9778.6142578, Throughput: 67.987427 im/s Epoch 42, Train Loss: 9565.6318359, Time: 73.4255s, Throughput: 67.987251 im/s Epoch 43, Iter 39, Loss: 9702.7597656, Throughput: 68.292909 im/s Epoch 43, Train Loss: 9527.6259766, Time: 73.0971s, Throughput: 68.292736 im/s Epoch 44, Iter 39, Loss: 9159.7744141, Throughput: 68.161749 im/s Epoch 44, Train Loss: 9563.1494141, Time: 73.2377s, Throughput: 68.161586 im/s Epoch 45, Iter 39, Loss: 9218.2128906, Throughput: 68.980294 im/s Epoch 45, Train Loss: 9473.7011719, Time: 72.3687s, Throughput: 68.980122 im/s Epoch 46, Iter 39, Loss: 9952.1992188, Throughput: 68.659445 im/s Epoch 46, Train Loss: 9378.1591797, Time: 72.7069s, Throughput: 68.659256 im/s Epoch 47, Iter 39, Loss: 9560.0312500, Throughput: 68.670714 im/s Epoch 47, Train Loss: 9346.6210938, Time: 72.6949s, Throughput: 68.670546 im/s Epoch 48, Iter 39, Loss: 9305.5390625, Throughput: 68.120295 im/s Epoch 48, Train Loss: 9318.7275391, Time: 73.2823s, Throughput: 68.120108 im/s Epoch 49, Iter 39, Loss: 9102.9833984, Throughput: 68.555589 im/s Epoch 49, Train Loss: 9292.0361328, Time: 72.8170s, Throughput: 68.555394 im/s Epoch 50, Iter 39, Loss: 9303.9628906, Throughput: 68.402689 im/s Epoch 50, Train Loss: 9338.2158203, Time: 72.9798s, Throughput: 68.402518 im/s ``` *** ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 7763 64-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver3) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/6_GCN_Cora.md --- # Graph Convolutional Networks on Cora {#GCN-Tutorial-Cora} This example is based on [GCN MLX tutorial](https://github.com/ml-explore/mlx-examples/blob/main/gcn/). While we are doing this manually, we recommend directly using [GNNLux.jl](https://juliagraphs.org/GraphNeuralNetworks.jl/docs/GNNLux.jl/stable/). ```julia using Lux, Reactant, MLDatasets, Random, Statistics, GNNGraphs, ConcreteStructs, Printf, OneHotArrays, Optimisers const xdev = reactant_device(; force=true) const cdev = cpu_device() ``` ## Loading Cora Dataset {#Loading-Cora-Dataset} ```julia function loadcora() data = Cora() gph = data.graphs[1] gnngraph = GNNGraph( gph.edge_index; ndata=gph.node_data, edata=gph.edge_data, gph.num_nodes ) return ( gph.node_data.features, onehotbatch(gph.node_data.targets, data.metadata["classes"]), # We use a dense matrix here to avoid incompatibility with Reactant Matrix{Int32}(adjacency_matrix(gnngraph)), # We use this since Reactant doesn't yet support gather adjoint (1:140, 141:640, 1709:2708), ) end ``` ## Model Definition {#Model-Definition} ```julia function GCNLayer(args...; kwargs...) return @compact(; dense=Dense(args...; kwargs...)) do (x, adj) @return dense(x) * adj end end function GCN(x_dim, h_dim, out_dim; nb_layers=2, dropout=0.5, kwargs...) layer_sizes = vcat(x_dim, [h_dim for _ in 1:nb_layers]) gcn_layers = [ GCNLayer(in_dim => out_dim; kwargs...) for (in_dim, out_dim) in zip(layer_sizes[1:(end - 1)], layer_sizes[2:end]) ] last_layer = GCNLayer(layer_sizes[end] => out_dim; kwargs...) dropout = Dropout(dropout) return @compact(; gcn_layers, dropout, last_layer) do (x, adj, mask) for layer in gcn_layers x = relu.(layer((x, adj))) x = dropout(x) end @return last_layer((x, adj))[:, mask] end end ``` ## Helper Functions {#Helper-Functions} ```julia function loss_function(model, ps, st, (x, y, adj, mask)) y_pred, st = model((x, adj, mask), ps, st) loss = CrossEntropyLoss(; agg=mean, logits=Val(true))(y_pred, y[:, mask]) return loss, st, (; y_pred) end accuracy(y_pred, y) = mean(onecold(y_pred) .== onecold(y)) * 100 ``` ## Training the Model {#Training-the-Model} ```julia function main(; hidden_dim::Int=64, dropout::Float64=0.1, nb_layers::Int=2, use_bias::Bool=true, lr::Float64=0.001, weight_decay::Float64=0.0, patience::Int=20, epochs::Int=200, ) rng = Random.default_rng() Random.seed!(rng, 0) features, targets, adj, (train_idx, val_idx, test_idx) = loadcora() |> xdev gcn = GCN(size(features, 1), hidden_dim, size(targets, 1); nb_layers, dropout, use_bias) ps, st = Lux.setup(rng, gcn) |> xdev opt = iszero(weight_decay) ? Adam(lr) : AdamW(; eta=lr, lambda=weight_decay) train_state = Training.TrainState(gcn, ps, st, opt) @printf "Total Trainable Parameters: %0.4f M\n" (Lux.parameterlength(ps) / 1.0e6) val_loss_compiled = @compile loss_function( gcn, ps, Lux.testmode(st), (features, targets, adj, val_idx) ) train_model_compiled = @compile gcn((features, adj, train_idx), ps, Lux.testmode(st)) val_model_compiled = @compile gcn((features, adj, val_idx), ps, Lux.testmode(st)) best_loss_val = Inf cnt = 0 for epoch in 1:epochs (_, loss, _, train_state) = Lux.Training.single_train_step!( AutoEnzyme(), loss_function, (features, targets, adj, train_idx), train_state; return_gradients=Val(false), ) train_acc = accuracy( Array( train_model_compiled( (features, adj, train_idx), train_state.parameters, Lux.testmode(train_state.states), )[1], ), Array(targets)[:, train_idx], ) val_loss = first( val_loss_compiled( gcn, train_state.parameters, Lux.testmode(train_state.states), (features, targets, adj, val_idx), ), ) val_acc = accuracy( Array( val_model_compiled( (features, adj, val_idx), train_state.parameters, Lux.testmode(train_state.states), )[1], ), Array(targets)[:, val_idx], ) @printf "Epoch %3d\tTrain Loss: %.6f\tTrain Acc: %.4f%%\tVal Loss: %.6f\t\ Val Acc: %.4f%%\n" epoch loss train_acc val_loss val_acc if val_loss < best_loss_val best_loss_val = val_loss cnt = 0 else cnt += 1 if cnt == patience @printf "Early Stopping at Epoch %d\n" epoch break end end end test_loss = @jit( loss_function( gcn, train_state.parameters, Lux.testmode(train_state.states), (features, targets, adj, test_idx), ) )[1] test_acc = accuracy( Array( @jit( gcn( (features, adj, test_idx), train_state.parameters, Lux.testmode(train_state.states), ) )[1], ), Array(targets)[:, test_idx], ) @printf "Test Loss: %.6f\tTest Acc: %.4f%%\n" test_loss test_acc return nothing end main() ``` ``` ┌ Warning: `replicate` doesn't work for `TaskLocalRNG`. Returning the same `TaskLocalRNG`. └ @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/LuxCore.jl:18 Total Trainable Parameters: 0.0964 M Epoch 1 Train Loss: 15.483308 Train Acc: 22.1429% Val Loss: 7.571783 Val Acc: 25.8000% Epoch 2 Train Loss: 10.125030 Train Acc: 22.1429% Val Loss: 3.797886 Val Acc: 29.4000% Epoch 3 Train Loss: 4.467242 Train Acc: 37.8571% Val Loss: 2.431701 Val Acc: 32.0000% Epoch 4 Train Loss: 2.424877 Train Acc: 51.4286% Val Loss: 2.113642 Val Acc: 37.8000% Epoch 5 Train Loss: 1.761382 Train Acc: 58.5714% Val Loss: 1.889251 Val Acc: 45.0000% Epoch 6 Train Loss: 1.484980 Train Acc: 67.8571% Val Loss: 1.611183 Val Acc: 51.6000% Epoch 7 Train Loss: 1.267712 Train Acc: 71.4286% Val Loss: 1.504884 Val Acc: 58.4000% Epoch 8 Train Loss: 1.319321 Train Acc: 72.1429% Val Loss: 1.505576 Val Acc: 59.8000% Epoch 9 Train Loss: 1.617086 Train Acc: 73.5714% Val Loss: 1.520861 Val Acc: 61.2000% Epoch 10 Train Loss: 1.249781 Train Acc: 74.2857% Val Loss: 1.519172 Val Acc: 62.0000% Epoch 11 Train Loss: 1.187690 Train Acc: 78.5714% Val Loss: 1.504537 Val Acc: 62.0000% Epoch 12 Train Loss: 1.179360 Train Acc: 78.5714% Val Loss: 1.547555 Val Acc: 61.8000% Epoch 13 Train Loss: 0.898748 Train Acc: 80.0000% Val Loss: 1.608347 Val Acc: 62.0000% Epoch 14 Train Loss: 0.946830 Train Acc: 80.0000% Val Loss: 1.649864 Val Acc: 61.8000% Epoch 15 Train Loss: 1.425961 Train Acc: 80.7143% Val Loss: 1.633293 Val Acc: 64.4000% Epoch 16 Train Loss: 0.875585 Train Acc: 82.1429% Val Loss: 1.616587 Val Acc: 66.6000% Epoch 17 Train Loss: 0.810615 Train Acc: 81.4286% Val Loss: 1.592887 Val Acc: 67.0000% Epoch 18 Train Loss: 0.763062 Train Acc: 80.7143% Val Loss: 1.569996 Val Acc: 67.4000% Epoch 19 Train Loss: 0.881348 Train Acc: 82.1429% Val Loss: 1.543069 Val Acc: 67.2000% Epoch 20 Train Loss: 0.750949 Train Acc: 82.8571% Val Loss: 1.520200 Val Acc: 66.8000% Epoch 21 Train Loss: 0.685395 Train Acc: 83.5714% Val Loss: 1.504100 Val Acc: 66.6000% Epoch 22 Train Loss: 0.611383 Train Acc: 85.0000% Val Loss: 1.500499 Val Acc: 66.0000% Epoch 23 Train Loss: 0.603166 Train Acc: 84.2857% Val Loss: 1.511355 Val Acc: 66.2000% Epoch 24 Train Loss: 1.565989 Train Acc: 85.7143% Val Loss: 1.550028 Val Acc: 66.0000% Epoch 25 Train Loss: 0.564262 Train Acc: 88.5714% Val Loss: 1.616222 Val Acc: 64.6000% Epoch 26 Train Loss: 0.524013 Train Acc: 87.8571% Val Loss: 1.695767 Val Acc: 64.0000% Epoch 27 Train Loss: 0.508034 Train Acc: 88.5714% Val Loss: 1.788846 Val Acc: 64.0000% Epoch 28 Train Loss: 0.621814 Train Acc: 87.8571% Val Loss: 1.853111 Val Acc: 63.0000% Epoch 29 Train Loss: 0.579144 Train Acc: 88.5714% Val Loss: 1.872775 Val Acc: 63.2000% Epoch 30 Train Loss: 0.491464 Train Acc: 88.5714% Val Loss: 1.874164 Val Acc: 63.8000% Epoch 31 Train Loss: 0.493936 Train Acc: 89.2857% Val Loss: 1.847677 Val Acc: 64.6000% Epoch 32 Train Loss: 0.562605 Train Acc: 90.0000% Val Loss: 1.800508 Val Acc: 66.0000% Epoch 33 Train Loss: 0.490371 Train Acc: 91.4286% Val Loss: 1.742706 Val Acc: 66.0000% Epoch 34 Train Loss: 0.623589 Train Acc: 91.4286% Val Loss: 1.702444 Val Acc: 65.8000% Epoch 35 Train Loss: 0.441532 Train Acc: 92.8571% Val Loss: 1.669238 Val Acc: 66.2000% Epoch 36 Train Loss: 0.414883 Train Acc: 92.1429% Val Loss: 1.649799 Val Acc: 67.4000% Epoch 37 Train Loss: 0.396852 Train Acc: 93.5714% Val Loss: 1.642260 Val Acc: 68.0000% Epoch 38 Train Loss: 0.370066 Train Acc: 93.5714% Val Loss: 1.644972 Val Acc: 68.2000% Epoch 39 Train Loss: 0.402366 Train Acc: 93.5714% Val Loss: 1.657053 Val Acc: 68.6000% Epoch 40 Train Loss: 0.802922 Train Acc: 95.7143% Val Loss: 1.677369 Val Acc: 67.8000% Epoch 41 Train Loss: 0.378652 Train Acc: 95.7143% Val Loss: 1.707681 Val Acc: 68.0000% Epoch 42 Train Loss: 0.366849 Train Acc: 95.0000% Val Loss: 1.735516 Val Acc: 68.2000% Early Stopping at Epoch 42 Test Loss: 1.518862 Test Acc: 68.8000% ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/7_RealNVP.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # Normalizing Flows for Density Estimation {#Normalizing-Flows-for-Density-Estimation} This tutorial demonstrates how to use Lux to train a [RealNVP](https://arxiv.org/abs/1605.08803). This is based on the [RealNVP implementation in MLX](https://github.com/ml-explore/mlx-examples/blob/main/normalizing_flow/). ```julia using Lux, Reactant, Random, Statistics, Enzyme, MLUtils, ConcreteStructs, Printf, Optimisers, CairoMakie const xdev = reactant_device(; force=true) const cdev = cpu_device() ``` ## Define & Load the Moons Dataset {#Define-and-Load-the-Moons-Dataset} We define a function to generate data from the moons dataset. We use the code here from [this tutorial](https://liorsinai.github.io/machine-learning/2024/08/19/micrograd-5-mlp.html#moons-dataset). ```julia function make_moons( rng::AbstractRNG, ::Type{T}, n_samples::Int=100; noise::Union{Nothing,AbstractFloat}=nothing, ) where {T} n_moons = n_samples ÷ 2 t_min, t_max = T(0), T(π) t_inner = rand(rng, T, n_moons) * (t_max - t_min) .+ t_min t_outer = rand(rng, T, n_moons) * (t_max - t_min) .+ t_min outer_circ_x = cos.(t_outer) outer_circ_y = sin.(t_outer) .+ T(1) inner_circ_x = 1 .- cos.(t_inner) inner_circ_y = 1 .- sin.(t_inner) .- T(1) data = [outer_circ_x outer_circ_y; inner_circ_x inner_circ_y] z = permutedims(data, (2, 1)) noise !== nothing && (z .+= T(noise) * randn(rng, T, size(z))) return z end ``` Let's visualize the dataset ```julia fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="x", ylabel="y") z = make_moons(Random.default_rng(), Float32, 10_000; noise=0.1) scatter!(ax, z[1, :], z[2, :]; markersize=2) fig ``` *** ```julia function load_moons_dataloader( args...; batchsize::Int, noise::Union{Nothing,AbstractFloat}=nothing, kwargs... ) return DataLoader( make_moons(args...; noise); batchsize, shuffle=true, partial=false, kwargs... ) end ``` ## Bijectors Implementation {#Bijectors-Implementation} ```julia abstract type AbstractBijector end @concrete struct AffineBijector <: AbstractBijector shift <: AbstractArray log_scale <: AbstractArray end function AffineBijector(shift_and_log_scale::AbstractArray{T,N}) where {T,N} n = size(shift_and_log_scale, 1) ÷ 2 idxs = ntuple(Returns(Colon()), N - 1) return AffineBijector( shift_and_log_scale[1:n, idxs...], shift_and_log_scale[(n + 1):end, idxs...] ) end function forward_and_log_det(bj::AffineBijector, x::AbstractArray) y = x .* exp.(bj.log_scale) .+ bj.shift return y, bj.log_scale end function inverse_and_log_det(bj::AffineBijector, y::AbstractArray) x = (y .- bj.shift) ./ exp.(bj.log_scale) return x, -bj.log_scale end @concrete struct MaskedCoupling <: AbstractBijector mask <: AbstractArray conditioner bijector end function apply_mask(bj::MaskedCoupling, x::AbstractArray, fn::F) where {F} x_masked = x .* (1 .- bj.mask) bijector_params = bj.conditioner(x_masked) y, log_det = fn(bijector_params) log_det = log_det .* bj.mask y = ifelse.(bj.mask, y, x) return y, dsum(log_det; dims=Tuple(collect(1:(ndims(x) - 1)))) end function forward_and_log_det(bj::MaskedCoupling, x::AbstractArray) return apply_mask(bj, x, params -> forward_and_log_det(bj.bijector(params), x)) end function inverse_and_log_det(bj::MaskedCoupling, y::AbstractArray) return apply_mask(bj, y, params -> inverse_and_log_det(bj.bijector(params), y)) end ``` ## Model Definition {#Model-Definition} ```julia function MLP(in_dims::Int, hidden_dims::Int, out_dims::Int, n_layers::Int; activation=gelu) return Chain( Dense(in_dims => hidden_dims, activation), [Dense(hidden_dims => hidden_dims, activation) for _ in 1:(n_layers - 1)]..., Dense(hidden_dims => out_dims), ) end @concrete struct RealNVP <: AbstractLuxContainerLayer{(:conditioners,)} conditioners dist_dims::Int n_transforms::Int end const StatefulRealNVP{M} = StatefulLuxLayer{M,<:RealNVP} function Lux.initialstates(rng::AbstractRNG, l::RealNVP) mask_list = Vector{Bool}[ collect(1:(l.dist_dims)) .% 2 .== i % 2 for i in 1:(l.n_transforms) ] return (; mask_list, conditioners=Lux.initialstates(rng, l.conditioners)) end function RealNVP(; n_transforms::Int, dist_dims::Int, hidden_dims::Int, n_layers::Int) conditioners = [ MLP(dist_dims, hidden_dims, 2 * dist_dims, n_layers; activation=gelu) for _ in 1:n_transforms ] conditioners = NamedTuple{ntuple(Base.Fix1(Symbol, :conditioners_), n_transforms)}( Tuple(conditioners) ) return RealNVP(conditioners, dist_dims, n_transforms) end log_prob(x::AbstractArray{T}) where {T} = -T(0.5 * log(2π)) .- T(0.5) .* abs2.(x) function log_prob(l::StatefulRealNVP, x::AbstractArray{T}) where {T} smodels = [ StatefulLuxLayer(conditioner, l.ps.conditioners[i], l.st.conditioners[i]) for (i, conditioner) in enumerate(l.model.conditioners) ] lprob = zeros_like(x, size(x, ndims(x))) for (mask, conditioner) in Iterators.reverse(zip(l.st.mask_list, smodels)) bj = MaskedCoupling(mask, conditioner, AffineBijector) x, log_det = inverse_and_log_det(bj, x) lprob += log_det end lprob += dsum(log_prob(x); dims=Tuple(collect(1:(ndims(x) - 1)))) conditioners = NamedTuple{ ntuple(Base.Fix1(Symbol, :conditioners_), l.model.n_transforms) }( Tuple([smodel.st for smodel in smodels]) ) l.st = merge(l.st, (; conditioners)) return lprob end function sample( rng::AbstractRNG, ::Type{T}, d::StatefulRealNVP, nsamples::Int, nsteps::Int=length(d.model.conditioners), ) where {T} @assert 1 ≤ nsteps ≤ length(d.model.conditioners) smodels = [ StatefulLuxLayer(conditioner, d.ps.conditioners[i], d.st.conditioners[i]) for (i, conditioner) in enumerate(d.model.conditioners) ] x = randn(rng, T, d.model.dist_dims, nsamples) for (i, (mask, conditioner)) in enumerate(zip(d.st.mask_list, smodels)) x, _ = forward_and_log_det(MaskedCoupling(mask, conditioner, AffineBijector), x) i ≥ nsteps && break end return x end ``` ## Helper Functions {#Helper-Functions} ```julia dsum(x; dims) = dropdims(sum(x; dims); dims) function loss_function(model, ps, st, x) smodel = StatefulLuxLayer(model, ps, st) lprob = log_prob(smodel, x) return -mean(lprob), smodel.st, (;) end ``` ## Training the Model {#Training-the-Model} ```julia function main(; maxiters::Int=10_000, n_train_samples::Int=100_000, batchsize::Int=128, n_transforms::Int=6, hidden_dims::Int=16, n_layers::Int=4, lr::Float64=0.0004, noise::Float64=0.06, ) rng = Random.default_rng() Random.seed!(rng, 0) dataloader = load_moons_dataloader(rng, Float32, n_train_samples; batchsize, noise) |> xdev |> Iterators.cycle model = RealNVP(; n_transforms, dist_dims=2, hidden_dims, n_layers) ps, st = Lux.setup(rng, model) |> xdev opt = Adam(lr) train_state = Training.TrainState(model, ps, st, opt) @printf "Total Trainable Parameters: %d\n" Lux.parameterlength(ps) total_samples = 0 start_time = time() for (iter, x) in enumerate(dataloader) total_samples += size(x, ndims(x)) (_, loss, _, train_state) = Training.single_train_step!( AutoEnzyme(), loss_function, x, train_state; return_gradients=Val(false) ) isnan(loss) && error("NaN loss encountered in iter $(iter)!") if iter == 1 || iter == maxiters || iter % 1000 == 0 throughput = total_samples / (time() - start_time) @printf "Iter: [%6d/%6d]\tTraining Loss: %.6f\t\ Throughput: %.6f samples/s\n" iter maxiters loss throughput end iter ≥ maxiters && break end return StatefulLuxLayer(model, train_state.parameters, Lux.testmode(train_state.states)) end trained_model = main() ``` ## Visualizing the Results {#Visualizing-the-Results} ```julia z_stages = Matrix{Float32}[] for i in 1:(trained_model.model.n_transforms) z = @jit sample(Random.default_rng(), Float32, trained_model, 10_000, i) push!(z_stages, Array(z)) end begin fig = Figure(; size=(1200, 800)) for (idx, z) in enumerate(z_stages) i, j = (idx - 1) ÷ 3, (idx - 1) % 3 ax = Axis(fig[i, j]; title="$(idx) transforms") scatter!(ax, z[1, :], z[2, :]; markersize=2) end fig end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/8_LSTMEncoderDecoder.md --- # Building a LSTM Encoder-Decoder model using Lux.jl {#Building-a-LSTM-Encoder-Decoder-model-using-Lux.jl} This examples is based on [LSTM\_encoder\_decoder](https://github.com/lkulowski/LSTM_encoder_decoder) by [Laura Kulowski](https://github.com/lkulowski). ```julia using Lux, Reactant, Random, Optimisers, Statistics, Enzyme, Printf, CairoMakie, MLUtils const xdev = reactant_device(; force=true) const cdev = cpu_device() ``` ## Generate synthetic data {#Generate-synthetic-data} ```julia function synthetic_data(Nt=2000, tf=80 * Float32(π)) t = range(0.0f0, tf; length=Nt) y = sin.(2.0f0 * t) .+ 0.5f0 * cos.(t) .+ randn(Float32, Nt) * 0.2f0 return t, y end function train_test_split(t, y, split=0.8) indx_split = ceil(Int, length(t) * split) indx_train = 1:indx_split indx_test = (indx_split + 1):length(t) t_train = t[indx_train] y_train = reshape(y[indx_train], 1, :) t_test = t[indx_test] y_test = reshape(y[indx_test], 1, :) return t_train, y_train, t_test, y_test end function windowed_dataset(y; input_window=5, output_window=1, stride=1, num_features=1) L = size(y, ndims(y)) num_samples = (L - input_window - output_window) ÷ stride + 1 X = zeros(Float32, num_features, input_window, num_samples) Y = zeros(Float32, num_features, output_window, num_samples) for ii in 1:num_samples, ff in 1:num_features start_x = stride * (ii - 1) + 1 end_x = start_x + input_window - 1 X[ff, :, ii] .= y[start_x:end_x] start_y = stride * (ii - 1) + input_window + 1 end_y = start_y + output_window - 1 Y[ff, :, ii] .= y[start_y:end_y] end return X, Y end t, y = synthetic_data() begin fig = Figure(; size=(1000, 400)) ax = Axis(fig[1, 1]; title="Synthetic Time Series", xlabel="t", ylabel="y") lines!(ax, t, y; label="y", color=:black, linewidth=2) fig end ``` *** ```julia t_train, y_train, t_test, y_test = train_test_split(t, y) begin fig = Figure(; size=(1000, 400)) ax = Axis( fig[1, 1]; title="Time Series Split into Train and Test Sets", xlabel="t", ylabel="y", ) lines!(ax, t_train, y_train[1, :]; label="Train", color=:black, linewidth=2) lines!(ax, t_test, y_test[1, :]; label="Test", color=:red, linewidth=2) fig[1, 2] = Legend(fig, ax) fig end ``` *** ```julia X_train, Y_train = windowed_dataset(y_train; input_window=80, output_window=20, stride=5) X_test, Y_test = windowed_dataset(y_test; input_window=80, output_window=20, stride=5) begin fig = Figure(; size=(1000, 400)) ax = Axis(fig[1, 1]; title="Example of Windowed Training Data", xlabel="t", ylabel="y") linestyles = [:solid, :dash, :dot, :dashdot, :dashdotdot] for b in 1:4:16 lines!( ax, 0:79, X_train[1, :, b]; label="Input", color=:black, linewidth=2, linestyle=linestyles[mod1(b, 5)], ) lines!( ax, 79:99, vcat(X_train[1, end, b], Y_train[1, :, b]); label="Target", color=:red, linewidth=2, linestyle=linestyles[mod1(b, 5)], ) end fig end ``` ## Define the model {#Define-the-model} ```julia struct RNNEncoder{C} <: AbstractLuxWrapperLayer{:cell} cell::C end function (rnn::RNNEncoder)(x::AbstractArray{T,3}, ps, st) where {T} (y, carry), st = Lux.apply(rnn.cell, x[:, 1, :], ps, st) @trace for i in 2:size(x, 2) (y, carry), st = Lux.apply(rnn.cell, (x[:, i, :], carry), ps, st) end return (y, carry), st end struct RNNDecoder{C,L} <: AbstractLuxContainerLayer{(:cell, :linear)} cell::C linear::L training_mode::Symbol teacher_forcing_ratio::Float32 function RNNDecoder( cell::C, linear::L; training_mode::Symbol=:recursive, teacher_forcing_ratio::Float32=0.5f0, ) where {C,L} @assert training_mode in (:recursive, :teacher_forcing, :mixed_teacher_forcing) return new{C,L}(cell, linear, training_mode, Float32(teacher_forcing_ratio)) end end function LuxCore.initialstates(rng::AbstractRNG, d::RNNDecoder) return (; cell=LuxCore.initialstates(rng, d.cell), linear=LuxCore.initialstates(rng, d.linear), training=Val(true), rng, ) end function _teacher_forcing_condition(::Val{false}, x, mode, rng, ratio, target_len) res = similar(x, Bool, target_len) fill!(res, true) return res end function _teacher_forcing_condition(::Val{true}, x, mode, rng, ratio, target_len) mode === :recursive && return _teacher_forcing_condition(Val(false), x, mode, rng, ratio, target_len) mode === :teacher_forcing && fill(rand(rng, Float32) < ratio, target_len) return rand(rng, Float32, target_len) .< ratio end function (rnn::RNNDecoder)((decoder_input, carry, target_len, target), ps, st) @assert ndims(decoder_input) == 2 rng = Lux.replicate(st.rng) if target === nothing ### This will be optimized out by Reactant target = similar( decoder_input, size(decoder_input, 1), target_len, size(decoder_input, 2) ) fill!(target, 0) else @assert size(target, 2) ≤ target_len end (y_latent, carry), st_cell = Lux.apply( rnn.cell, (decoder_input, carry), ps.cell, st.cell ) y_pred, st_linear = Lux.apply(rnn.linear, y_latent, ps.linear, st.linear) y_full = similar(y_pred, size(y_pred, 1), target_len, size(y_pred, 2)) y_full[:, 1, :] = y_pred conditions = _teacher_forcing_condition( st.training, decoder_input, rnn.training_mode, rng, rnn.teacher_forcing_ratio, target_len, ) decoder_input = ifelse.(@allowscalar(conditions[1]), target[:, 1, :], y_pred) @trace for i in 2:target_len (y_latent, carry), st_cell = Lux.apply( rnn.cell, (decoder_input, carry), ps.cell, st_cell ) y_pred, st_linear = Lux.apply(rnn.linear, y_latent, ps.linear, st_linear) y_full[:, i, :] = y_pred decoder_input = ifelse.(@allowscalar(conditions[i]), target[:, i, :], y_pred) end return y_full, merge(st, (; cell=st_cell, linear=st_linear, rng)) end struct RNNEncoderDecoder{C<:RNNEncoder,L<:RNNDecoder} <: AbstractLuxContainerLayer{(:encoder, :decoder)} encoder::C decoder::L end function (rnn::RNNEncoderDecoder)((x, target_len, target), ps, st) (y, carry), st_encoder = Lux.apply(rnn.encoder, x, ps.encoder, st.encoder) pred, st_decoder = Lux.apply( rnn.decoder, (x[:, end, :], carry, target_len, target), ps.decoder, st.decoder ) return pred, (; encoder=st_encoder, decoder=st_decoder) end ``` ## Training {#Training} ```julia function train( train_dataset, validation_dataset; nepochs=50, batchsize=32, hidden_dims=32, training_mode=:mixed_teacher_forcing, teacher_forcing_ratio=0.5f0, learning_rate=1e-3, ) (X_train, Y_train), (X_test, Y_test) = train_dataset, validation_dataset in_dims = size(X_train, 1) @assert size(Y_train, 2) == size(Y_test, 2) target_len = size(Y_train, 2) train_dataloader = DataLoader( (X_train, Y_train); batchsize=min(batchsize, size(X_train, 4)), shuffle=true, partial=false, ) |> xdev X_test, Y_test = (X_test, Y_test) |> xdev model = RNNEncoderDecoder( RNNEncoder(LSTMCell(in_dims => hidden_dims)), RNNDecoder( LSTMCell(in_dims => hidden_dims), Dense(hidden_dims => in_dims); training_mode, teacher_forcing_ratio, ), ) ps, st = Lux.setup(Random.default_rng(), model) |> xdev train_state = Training.TrainState(model, ps, st, Optimisers.Adam(learning_rate)) stime = time() model_compiled = @compile model((X_test, target_len, nothing), ps, Lux.testmode(st)) ttime = time() - stime @printf "Compilation time: %.4f seconds\n\n" ttime for epoch in 1:nepochs stime = time() for (x, y) in train_dataloader (_, _, _, train_state) = Training.single_train_step!( AutoEnzyme(), MSELoss(), ((x, target_len, y), y), train_state; return_gradients=Val(false), ) end ttime = time() - stime y_pred, _ = model_compiled( (X_test, target_len, nothing), train_state.parameters, Lux.testmode(train_state.states), ) pred_loss = Float32(@jit(MSELoss()(y_pred, Y_test))) @printf( "[%3d/%3d]\tTime per epoch: %3.5fs\tValidation Loss: %.4f\n", epoch, nepochs, ttime, pred_loss, ) end return StatefulLuxLayer( model, train_state.parameters |> cdev, train_state.states |> cdev ) end trained_model = train( (X_train, Y_train), (X_test, Y_test); nepochs=50, batchsize=4, hidden_dims=32, training_mode=:mixed_teacher_forcing, teacher_forcing_ratio=0.5f0, learning_rate=3e-4, ) ``` ``` StatefulLuxLayer{Val{true}()}( RNNEncoderDecoder( encoder = RNNEncoder( cell = LSTMCell(1 => 32), # 4_480 parameters, plus 1 non-trainable ), decoder = RNNDecoder( cell = LSTMCell(1 => 32), # 4_480 parameters, plus 1 non-trainable linear = Dense(32 => 1), # 33 parameters ), ), ) # Total: 8_993 parameters, # plus 2 states. ``` ## Making Predictions {#Making-Predictions} ```julia Y_pred = trained_model((X_test, 20, nothing)) begin fig = Figure(; size=(1200, 800)) for i in 1:4, j in 1:2 b = i + j * 4 ax = Axis(fig[i, j]; xlabel="t", ylabel="y") i != 4 && hidexdecorations!(ax; grid=false) j != 1 && hideydecorations!(ax; grid=false) lines!(ax, 0:79, X_test[1, :, b]; label="Input", color=:black, linewidth=2) lines!( ax, 79:99, vcat(X_test[1, end, b], Y_test[1, :, b]); label="Ground Truth\n(Noisy)", color=:red, linewidth=2, ) lines!( ax, 79:99, vcat(X_test[1, end, b], Y_pred[1, :, b]); label="Prediction", color=:blue, linewidth=2, ) i == 4 && j == 2 && axislegend(ax; position=:lb) end fig[0, :] = Label(fig, "Predictions from Trained Model"; fontsize=20) fig end ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/9_CIFAR10_conv_mixer.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # ConvMixer on CIFAR-10 {#ConvMixer-on-CIFAR-10} ## Package Imports {#Package-Imports} ```julia using Comonicon, Interpolations, Lux, Optimisers, Printf, Random, Statistics ``` Set some global flags that will improve performance ```julia XLA_FLAGS = get(ENV, "XLA_FLAGS", "") ENV["XLA_FLAGS"] = "$(XLA_FLAGS) --xla_gpu_enable_cublaslt=true" ``` ## Load Common Packages {#Load-Common-Packages} ```julia using ConcreteStructs, DataAugmentation, ImageShow, Lux, MLDatasets, MLUtils, OneHotArrays, Printf, ProgressTables, Random, BFloat16s using Reactant ``` ## Data Loading Functionality {#Data-Loading-Functionality} ```julia @concrete struct TensorDataset dataset transform end Base.length(ds::TensorDataset) = length(ds.dataset) function Base.getindex(ds::TensorDataset, idxs::Union{Vector{<:Integer},AbstractRange}) img = Image.(eachslice(convert2image(ds.dataset, idxs); dims=3)) y = onehotbatch(ds.dataset.targets[idxs], 0:9) return stack(parent ∘ itemdata ∘ Base.Fix1(apply, ds.transform), img), y end function get_cifar10_dataloaders(::Type{T}, batchsize; kwargs...) where {T} cifar10_mean = T.((0.4914, 0.4822, 0.4465)) cifar10_std = T.((0.2471, 0.2435, 0.2616)) train_transform = RandomResizeCrop((32, 32)) |> Maybe(FlipX{2}()) |> ImageToTensor() |> Normalize(cifar10_mean, cifar10_std) |> ToEltype(T) test_transform = ImageToTensor() |> Normalize(cifar10_mean, cifar10_std) |> ToEltype(T) trainset = TensorDataset(CIFAR10(; Tx=T, split=:train), train_transform) trainloader = DataLoader(trainset; batchsize, shuffle=true, kwargs...) testset = TensorDataset(CIFAR10(; Tx=T, split=:test), test_transform) testloader = DataLoader(testset; batchsize, shuffle=false, kwargs...) return trainloader, testloader end ``` ## Utility Functions {#Utility-Functions} ```julia function accuracy(model, ps, st, dataloader) total_correct, total = 0, 0 cdev = cpu_device() for (x, y) in dataloader target_class = onecold(cdev(y)) predicted_class = onecold(cdev(first(model(x, ps, st)))) total_correct += sum(target_class .== predicted_class) total += length(target_class) end return total_correct / total end ``` ## Training Loop {#Training-Loop} ```julia function train_model( model, opt, scheduler=nothing; batchsize::Int=512, seed::Int=1234, epochs::Int=25, bfloat16::Bool=false, ) rng = Random.default_rng() Random.seed!(rng, seed) prec = bfloat16 ? bf16 : f32 prec_jl = bfloat16 ? BFloat16 : Float32 prec_str = bfloat16 ? "BFloat16" : "Float32" @printf "[Info] Using %s precision\n" prec_str dev = reactant_device(; force=true) trainloader, testloader = get_cifar10_dataloaders(prec_jl, batchsize; partial=false) |> dev ps, st = prec(Lux.setup(rng, model)) |> dev train_state = Training.TrainState(model, ps, st, opt) x_ra = rand(rng, prec_jl, size(first(trainloader)[1])) |> dev @printf "[Info] Compiling model with Reactant.jl\n" model_compiled = @compile model(x_ra, ps, Lux.testmode(st)) @printf "[Info] Model compiled!\n" loss_fn = CrossEntropyLoss(; logits=Val(true)) pt = ProgressTable(; header=[ "Epoch", "Learning Rate", "Train Accuracy (%)", "Test Accuracy (%)", "Time (s)" ], widths=[24, 24, 24, 24, 24], format=["%3d", "%.6f", "%.6f", "%.6f", "%.6f"], color=[:normal, :normal, :blue, :blue, :normal], border=true, alignment=[:center, :center, :center, :center, :center], ) @printf "[Info] Training model\n" initialize(pt) for epoch in 1:epochs stime = time() lr = 0 for (i, (x, y)) in enumerate(trainloader) if scheduler !== nothing lr = scheduler((epoch - 1) + (i + 1) / length(trainloader)) train_state = Optimisers.adjust!(train_state, lr) end (_, loss, _, train_state) = Training.single_train_step!( AutoEnzyme(), loss_fn, (x, y), train_state; return_gradients=Val(false) ) isnan(loss) && error("NaN loss encountered!") end ttime = time() - stime train_acc = accuracy( model_compiled, train_state.parameters, Lux.testmode(train_state.states), trainloader, ) * 100 test_acc = accuracy( model_compiled, train_state.parameters, Lux.testmode(train_state.states), testloader, ) * 100 scheduler === nothing && (lr = NaN32) next(pt, [epoch, lr, train_acc, test_acc, ttime]) end finalize(pt) return @printf "[Info] Finished training\n" end ``` ## Model Definition {#Model-Definition} ```julia function ConvMixer(; dim, depth, kernel_size=5, patch_size=2) return Chain( Conv((patch_size, patch_size), 3 => dim, relu; stride=patch_size), BatchNorm(dim), [ Chain( SkipConnection( Chain( Conv( (kernel_size, kernel_size), dim => dim, relu; groups=dim, pad=SamePad(), ), BatchNorm(dim), ), +, ), Conv((1, 1), dim => dim, relu), BatchNorm(dim), ) for _ in 1:depth ]..., GlobalMeanPool(), FlattenLayer(), Dense(dim => 10), ) end ``` ## Entry Point {#Entry-Point} ```julia Comonicon.@main function main(; batchsize::Int=512, hidden_dim::Int=256, depth::Int=8, patch_size::Int=2, kernel_size::Int=5, weight_decay::Float64=0.0001, clip_norm::Bool=false, seed::Int=1234, epochs::Int=25, lr_max::Float64=0.05, bfloat16::Bool=false, ) model = ConvMixer(; dim=hidden_dim, depth, kernel_size, patch_size) opt = AdamW(; eta=lr_max, lambda=weight_decay) clip_norm && (opt = OptimiserChain(ClipNorm(), opt)) lr_schedule = linear_interpolation( [0, epochs * 2 ÷ 5, epochs * 4 ÷ 5, epochs + 1], [0, lr_max, lr_max / 20, 0] ) return train_model(model, opt, lr_schedule; batchsize, seed, epochs, bfloat16) end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/10_CIFAR10_simple_cnn.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # Convolutional Neural Network on CIFAR-10 {#Convolutional-Neural-Network-on-CIFAR-10} ## Package Imports {#Package-Imports} ```julia using Comonicon, Lux, Optimisers, Printf, Random, Statistics, Enzyme ``` Set some global flags that will improve performance ```julia XLA_FLAGS = get(ENV, "XLA_FLAGS", "") ENV["XLA_FLAGS"] = "$(XLA_FLAGS) --xla_gpu_enable_cublaslt=true" ``` ## Load Common Packages {#Load-Common-Packages} ```julia using ConcreteStructs, DataAugmentation, ImageShow, Lux, MLDatasets, MLUtils, OneHotArrays, Printf, ProgressTables, Random, BFloat16s using Reactant ``` ## Data Loading Functionality {#Data-Loading-Functionality} ```julia @concrete struct TensorDataset dataset transform end Base.length(ds::TensorDataset) = length(ds.dataset) function Base.getindex(ds::TensorDataset, idxs::Union{Vector{<:Integer},AbstractRange}) img = Image.(eachslice(convert2image(ds.dataset, idxs); dims=3)) y = onehotbatch(ds.dataset.targets[idxs], 0:9) return stack(parent ∘ itemdata ∘ Base.Fix1(apply, ds.transform), img), y end function get_cifar10_dataloaders(::Type{T}, batchsize; kwargs...) where {T} cifar10_mean = T.((0.4914, 0.4822, 0.4465)) cifar10_std = T.((0.2471, 0.2435, 0.2616)) train_transform = RandomResizeCrop((32, 32)) |> Maybe(FlipX{2}()) |> ImageToTensor() |> Normalize(cifar10_mean, cifar10_std) |> ToEltype(T) test_transform = ImageToTensor() |> Normalize(cifar10_mean, cifar10_std) |> ToEltype(T) trainset = TensorDataset(CIFAR10(; Tx=T, split=:train), train_transform) trainloader = DataLoader(trainset; batchsize, shuffle=true, kwargs...) testset = TensorDataset(CIFAR10(; Tx=T, split=:test), test_transform) testloader = DataLoader(testset; batchsize, shuffle=false, kwargs...) return trainloader, testloader end ``` ## Utility Functions {#Utility-Functions} ```julia function accuracy(model, ps, st, dataloader) total_correct, total = 0, 0 cdev = cpu_device() for (x, y) in dataloader target_class = onecold(cdev(y)) predicted_class = onecold(cdev(first(model(x, ps, st)))) total_correct += sum(target_class .== predicted_class) total += length(target_class) end return total_correct / total end ``` ## Training Loop {#Training-Loop} ```julia function train_model( model, opt, scheduler=nothing; batchsize::Int=512, seed::Int=1234, epochs::Int=25, bfloat16::Bool=false, ) rng = Random.default_rng() Random.seed!(rng, seed) prec = bfloat16 ? bf16 : f32 prec_jl = bfloat16 ? BFloat16 : Float32 prec_str = bfloat16 ? "BFloat16" : "Float32" @printf "[Info] Using %s precision\n" prec_str dev = reactant_device(; force=true) trainloader, testloader = get_cifar10_dataloaders(prec_jl, batchsize; partial=false) |> dev ps, st = prec(Lux.setup(rng, model)) |> dev train_state = Training.TrainState(model, ps, st, opt) x_ra = rand(rng, prec_jl, size(first(trainloader)[1])) |> dev @printf "[Info] Compiling model with Reactant.jl\n" model_compiled = @compile model(x_ra, ps, Lux.testmode(st)) @printf "[Info] Model compiled!\n" loss_fn = CrossEntropyLoss(; logits=Val(true)) pt = ProgressTable(; header=[ "Epoch", "Learning Rate", "Train Accuracy (%)", "Test Accuracy (%)", "Time (s)" ], widths=[24, 24, 24, 24, 24], format=["%3d", "%.6f", "%.6f", "%.6f", "%.6f"], color=[:normal, :normal, :blue, :blue, :normal], border=true, alignment=[:center, :center, :center, :center, :center], ) @printf "[Info] Training model\n" initialize(pt) for epoch in 1:epochs stime = time() lr = 0 for (i, (x, y)) in enumerate(trainloader) if scheduler !== nothing lr = scheduler((epoch - 1) + (i + 1) / length(trainloader)) train_state = Optimisers.adjust!(train_state, lr) end (_, loss, _, train_state) = Training.single_train_step!( AutoEnzyme(), loss_fn, (x, y), train_state; return_gradients=Val(false) ) isnan(loss) && error("NaN loss encountered!") end ttime = time() - stime train_acc = accuracy( model_compiled, train_state.parameters, Lux.testmode(train_state.states), trainloader, ) * 100 test_acc = accuracy( model_compiled, train_state.parameters, Lux.testmode(train_state.states), testloader, ) * 100 scheduler === nothing && (lr = NaN32) next(pt, [epoch, lr, train_acc, test_acc, ttime]) end finalize(pt) return @printf "[Info] Finished training\n" end ``` ## Model Definition {#Model-Definition} ```julia function SimpleCNN() return Chain( Chain( Conv((3, 3), 3 => 16, relu; stride=2, pad=1), BatchNorm(16), Conv((3, 3), 16 => 32, relu; stride=2, pad=1), BatchNorm(32), Conv((3, 3), 32 => 64, relu; stride=2, pad=1), BatchNorm(64), Conv((3, 3), 64 => 128, relu; stride=2, pad=1), BatchNorm(128), ), GlobalMeanPool(), FlattenLayer(), Chain(Dense(128 => 64, relu), BatchNorm(64), Dense(64 => 10)), ) end ``` ## Entry Point {#Entry-Point} ```julia Comonicon.@main function main(; batchsize::Int=512, weight_decay::Float64=0.0001, clip_norm::Bool=false, seed::Int=1234, epochs::Int=50, lr::Float64=0.003, bfloat16::Bool=false, ) model = SimpleCNN() opt = AdamW(; eta=lr, lambda=weight_decay) clip_norm && (opt = OptimiserChain(ClipNorm(), opt)) return train_model(model, opt, nothing; batchsize, seed, epochs, bfloat16) end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/intermediate/11_CIFAR10_resnet20.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # ResNet20 on CIFAR-10 {#ResNet20-on-CIFAR-10} ## Package Imports {#Package-Imports} ```julia using Comonicon, Lux, Optimisers, Printf, Random, Statistics ``` Set some global flags that will improve performance ```julia XLA_FLAGS = get(ENV, "XLA_FLAGS", "") ENV["XLA_FLAGS"] = "$(XLA_FLAGS) --xla_gpu_enable_cublaslt=true" ``` ## Load Common Packages {#Load-Common-Packages} ```julia using ConcreteStructs, DataAugmentation, ImageShow, Lux, MLDatasets, MLUtils, OneHotArrays, Printf, ProgressTables, Random, BFloat16s using Reactant ``` ## Data Loading Functionality {#Data-Loading-Functionality} ```julia @concrete struct TensorDataset dataset transform end Base.length(ds::TensorDataset) = length(ds.dataset) function Base.getindex(ds::TensorDataset, idxs::Union{Vector{<:Integer},AbstractRange}) img = Image.(eachslice(convert2image(ds.dataset, idxs); dims=3)) y = onehotbatch(ds.dataset.targets[idxs], 0:9) return stack(parent ∘ itemdata ∘ Base.Fix1(apply, ds.transform), img), y end function get_cifar10_dataloaders(::Type{T}, batchsize; kwargs...) where {T} cifar10_mean = T.((0.4914, 0.4822, 0.4465)) cifar10_std = T.((0.2471, 0.2435, 0.2616)) train_transform = RandomResizeCrop((32, 32)) |> Maybe(FlipX{2}()) |> ImageToTensor() |> Normalize(cifar10_mean, cifar10_std) |> ToEltype(T) test_transform = ImageToTensor() |> Normalize(cifar10_mean, cifar10_std) |> ToEltype(T) trainset = TensorDataset(CIFAR10(; Tx=T, split=:train), train_transform) trainloader = DataLoader(trainset; batchsize, shuffle=true, kwargs...) testset = TensorDataset(CIFAR10(; Tx=T, split=:test), test_transform) testloader = DataLoader(testset; batchsize, shuffle=false, kwargs...) return trainloader, testloader end ``` ## Utility Functions {#Utility-Functions} ```julia function accuracy(model, ps, st, dataloader) total_correct, total = 0, 0 cdev = cpu_device() for (x, y) in dataloader target_class = onecold(cdev(y)) predicted_class = onecold(cdev(first(model(x, ps, st)))) total_correct += sum(target_class .== predicted_class) total += length(target_class) end return total_correct / total end ``` ## Training Loop {#Training-Loop} ```julia function train_model( model, opt, scheduler=nothing; batchsize::Int=512, seed::Int=1234, epochs::Int=25, bfloat16::Bool=false, ) rng = Random.default_rng() Random.seed!(rng, seed) prec = bfloat16 ? bf16 : f32 prec_jl = bfloat16 ? BFloat16 : Float32 prec_str = bfloat16 ? "BFloat16" : "Float32" @printf "[Info] Using %s precision\n" prec_str dev = reactant_device(; force=true) trainloader, testloader = get_cifar10_dataloaders(prec_jl, batchsize; partial=false) |> dev ps, st = prec(Lux.setup(rng, model)) |> dev train_state = Training.TrainState(model, ps, st, opt) x_ra = rand(rng, prec_jl, size(first(trainloader)[1])) |> dev @printf "[Info] Compiling model with Reactant.jl\n" model_compiled = @compile model(x_ra, ps, Lux.testmode(st)) @printf "[Info] Model compiled!\n" loss_fn = CrossEntropyLoss(; logits=Val(true)) pt = ProgressTable(; header=[ "Epoch", "Learning Rate", "Train Accuracy (%)", "Test Accuracy (%)", "Time (s)" ], widths=[24, 24, 24, 24, 24], format=["%3d", "%.6f", "%.6f", "%.6f", "%.6f"], color=[:normal, :normal, :blue, :blue, :normal], border=true, alignment=[:center, :center, :center, :center, :center], ) @printf "[Info] Training model\n" initialize(pt) for epoch in 1:epochs stime = time() lr = 0 for (i, (x, y)) in enumerate(trainloader) if scheduler !== nothing lr = scheduler((epoch - 1) + (i + 1) / length(trainloader)) train_state = Optimisers.adjust!(train_state, lr) end (_, loss, _, train_state) = Training.single_train_step!( AutoEnzyme(), loss_fn, (x, y), train_state; return_gradients=Val(false) ) isnan(loss) && error("NaN loss encountered!") end ttime = time() - stime train_acc = accuracy( model_compiled, train_state.parameters, Lux.testmode(train_state.states), trainloader, ) * 100 test_acc = accuracy( model_compiled, train_state.parameters, Lux.testmode(train_state.states), testloader, ) * 100 scheduler === nothing && (lr = NaN32) next(pt, [epoch, lr, train_acc, test_acc, ttime]) end finalize(pt) return @printf "[Info] Finished training\n" end ``` ## Model Definition {#Model-Definition} ```julia function ConvBN(kernel_size, (in_chs, out_chs), act; kwargs...) return Chain(Conv(kernel_size, in_chs => out_chs, act; kwargs...), BatchNorm(out_chs)) end function BasicBlock(in_channels, out_channels; stride=1) connection = if (stride == 1 && in_channels == out_channels) NoOpLayer() else Conv((3, 3), in_channels => out_channels, identity; stride=stride, pad=SamePad()) end return Chain( Parallel( +, connection, Chain( ConvBN((3, 3), in_channels => out_channels, relu; stride, pad=SamePad()), ConvBN((3, 3), out_channels => out_channels, identity; pad=SamePad()), ), ), Base.BroadcastFunction(relu), ) end function ResNet20(; num_classes=10) layers = [] # Initial Conv Layer push!(layers, Chain(Conv((3, 3), 3 => 16, relu; pad=SamePad()), BatchNorm(16))) # Residual Blocks block_configs = [ # (in_channels, out_channels, num_blocks, stride) (16, 16, 3, 1), (16, 32, 3, 2), (32, 64, 3, 2), ] for (in_channels, out_channels, num_blocks, stride) in block_configs for i in 1:num_blocks push!( layers, BasicBlock( i == 1 ? in_channels : out_channels, out_channels; stride=(i == 1 ? stride : 1), ), ) end end # Global Pooling and Final Dense Layer push!(layers, GlobalMeanPool()) push!(layers, FlattenLayer()) push!(layers, Dense(64 => num_classes)) return Chain(layers...) end ``` ## Entry Point {#Entry-Point} ```julia Comonicon.@main function main(; batchsize::Int=512, weight_decay::Float64=0.0001, clip_norm::Bool=false, seed::Int=1234, epochs::Int=100, lr::Float64=0.001, bfloat16::Bool=false, ) model = ResNet20() opt = AdamW(; eta=lr, lambda=weight_decay) clip_norm && (opt = OptimiserChain(ClipNorm(), opt)) lr_schedule = nothing return train_model(model, opt, lr_schedule; batchsize, seed, epochs, bfloat16) end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/advanced/1_GravitationalWaveForm.md --- # Training a Neural ODE to Model Gravitational Waveforms {#Training-a-Neural-ODE-to-Model-Gravitational-Waveforms} This code is adapted from [Astroinformatics/ScientificMachineLearning](https://github.com/Astroinformatics/ScientificMachineLearning/blob/c93aac3a460d70b4cce98836b677fd9b732e94b7/neuralode_gw.ipynb) The code has been minimally adapted from [Keith et. al. 2021](https://arxiv.org/abs/2102.12695) which originally used Flux.jl ## Package Imports {#Package-Imports} ```julia using Lux, ComponentArrays, LineSearches, OrdinaryDiffEqLowOrderRK, Optimization, OptimizationOptimJL, Printf, Random, SciMLSensitivity using CairoMakie ``` ## Define some Utility Functions {#Define-some-Utility-Functions} ::: tip Tip This section can be skipped. It defines functions to simulate the model, however, from a scientific machine learning perspective, isn't super relevant. ::: We need a very crude 2-body path. Assume the 1-body motion is a Newtonian 2-body position vector $r = r\_1 - r\_2$ and use Newtonian formulas to get $r\_1$, $r\_2$ (e.g. Theoretical Mechanics of Particles and Continua 4.3) ```julia function one2two(path, m₁, m₂) M = m₁ + m₂ r₁ = m₂ / M .* path r₂ = -m₁ / M .* path return r₁, r₂ end ``` Next we define a function to perform the change of variables: $(\chi(t),\phi(t)) \mapsto (x(t),y(t))$ ```julia @views function soln2orbit(soln, model_params=nothing) @assert size(soln, 1) ∈ [2, 4] "size(soln,1) must be either 2 or 4" if size(soln, 1) == 2 χ = soln[1, :] ϕ = soln[2, :] @assert length(model_params) == 3 "model_params must have length 3 when size(soln,2) = 2" p, M, e = model_params else χ = soln[1, :] ϕ = soln[2, :] p = soln[3, :] e = soln[4, :] end r = p ./ (1 .+ e .* cos.(χ)) x = r .* cos.(ϕ) y = r .* sin.(ϕ) orbit = vcat(x', y') return orbit end ``` This function uses second-order one-sided difference stencils at the endpoints; see https://doi.org/10.1090/S0025-5718-1988-0935077-0 ```julia function d_dt(v::AbstractVector, dt) a = -3 / 2 * v[1] + 2 * v[2] - 1 / 2 * v[3] b = (v[3:end] .- v[1:(end - 2)]) / 2 c = 3 / 2 * v[end] - 2 * v[end - 1] + 1 / 2 * v[end - 2] return [a; b; c] / dt end ``` This function uses second-order one-sided difference stencils at the endpoints; see https://doi.org/10.1090/S0025-5718-1988-0935077-0 ```julia function d2_dt2(v::AbstractVector, dt) a = 2 * v[1] - 5 * v[2] + 4 * v[3] - v[4] b = v[1:(end - 2)] .- 2 * v[2:(end - 1)] .+ v[3:end] c = 2 * v[end] - 5 * v[end - 1] + 4 * v[end - 2] - v[end - 3] return [a; b; c] / (dt^2) end ``` Now we define a function to compute the trace-free moment tensor from the orbit ```julia function orbit2tensor(orbit, component, mass=1) x = orbit[1, :] y = orbit[2, :] Ixx = x .^ 2 Iyy = y .^ 2 Ixy = x .* y trace = Ixx .+ Iyy if component[1] == 1 && component[2] == 1 tmp = Ixx .- trace ./ 3 elseif component[1] == 2 && component[2] == 2 tmp = Iyy .- trace ./ 3 else tmp = Ixy end return mass .* tmp end function h_22_quadrupole_components(dt, orbit, component, mass=1) mtensor = orbit2tensor(orbit, component, mass) mtensor_ddot = d2_dt2(mtensor, dt) return 2 * mtensor_ddot end function h_22_quadrupole(dt, orbit, mass=1) h11 = h_22_quadrupole_components(dt, orbit, (1, 1), mass) h22 = h_22_quadrupole_components(dt, orbit, (2, 2), mass) h12 = h_22_quadrupole_components(dt, orbit, (1, 2), mass) return h11, h12, h22 end function h_22_strain_one_body(dt::T, orbit) where {T} h11, h12, h22 = h_22_quadrupole(dt, orbit) h₊ = h11 - h22 hₓ = T(2) * h12 scaling_const = √(T(π) / 5) return scaling_const * h₊, -scaling_const * hₓ end function h_22_quadrupole_two_body(dt, orbit1, mass1, orbit2, mass2) h11_1, h12_1, h22_1 = h_22_quadrupole(dt, orbit1, mass1) h11_2, h12_2, h22_2 = h_22_quadrupole(dt, orbit2, mass2) h11 = h11_1 + h11_2 h12 = h12_1 + h12_2 h22 = h22_1 + h22_2 return h11, h12, h22 end function h_22_strain_two_body(dt::T, orbit1, mass1, orbit2, mass2) where {T} # compute (2,2) mode strain from orbits of BH1 of mass1 and BH2 of mass 2 @assert abs(mass1 + mass2 - 1.0) < 1.0e-12 "Masses do not sum to unity" h11, h12, h22 = h_22_quadrupole_two_body(dt, orbit1, mass1, orbit2, mass2) h₊ = h11 - h22 hₓ = T(2) * h12 scaling_const = √(T(π) / 5) return scaling_const * h₊, -scaling_const * hₓ end function compute_waveform(dt::T, soln, mass_ratio, model_params=nothing) where {T} @assert mass_ratio ≤ 1 "mass_ratio must be <= 1" @assert mass_ratio ≥ 0 "mass_ratio must be non-negative" orbit = soln2orbit(soln, model_params) if mass_ratio > 0 m₂ = inv(T(1) + mass_ratio) m₁ = mass_ratio * m₂ orbit₁, orbit₂ = one2two(orbit, m₁, m₂) waveform = h_22_strain_two_body(dt, orbit₁, m₁, orbit₂, m₂) else waveform = h_22_strain_one_body(dt, orbit) end return waveform end ``` ## Simulating the True Model {#Simulating-the-True-Model} `RelativisticOrbitModel` defines system of odes which describes motion of point like particle in Schwarzschild background, uses $$u\[1] = \chi$$ $$u\[2] = \phi$$ where, $p$, $M$, and $e$ are constants ```julia function RelativisticOrbitModel(u, (p, M, e), t) χ, ϕ = u numer = (p - 2 - 2 * e * cos(χ)) * (1 + e * cos(χ))^2 denom = sqrt((p - 2)^2 - 4 * e^2) χ̇ = numer * sqrt(p - 6 - 2 * e * cos(χ)) / (M * (p^2) * denom) ϕ̇ = numer / (M * (p^(3 / 2)) * denom) return [χ̇, ϕ̇] end mass_ratio = 0.0 # test particle u0 = Float64[π, 0.0] # initial conditions datasize = 250 tspan = (0.0f0, 6.0f4) # timespace for GW waveform tsteps = range(tspan[1], tspan[2]; length=datasize) # time at each timestep dt_data = tsteps[2] - tsteps[1] dt = 100.0 const ode_model_params = [100.0, 1.0, 0.5]; # p, M, e ``` Let's simulate the true model and plot the results using `OrdinaryDiffEq.jl` ```julia prob = ODEProblem(RelativisticOrbitModel, u0, tspan, ode_model_params) soln = Array(solve(prob, RK4(); saveat=tsteps, dt, adaptive=false)) waveform = first(compute_waveform(dt_data, soln, mass_ratio, ode_model_params)) begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="Time", ylabel="Waveform") l = lines!(ax, tsteps, waveform; linewidth=2, alpha=0.75) s = scatter!(ax, tsteps, waveform; marker=:circle, markersize=12, alpha=0.5) axislegend(ax, [[l, s]], ["Waveform Data"]) fig end ``` ## Defining a Neural Network Model {#Defining-a-Neural-Network-Model} Next, we define the neural network model that takes 1 input (time) and has two outputs. We'll make a function `ODE_model` that takes the initial conditions, neural network parameters and a time as inputs and returns the derivatives. It is typically never recommended to use globals but in case you do use them, make sure to mark them as `const`. We will deviate from the standard Neural Network initialization and use `WeightInitializers.jl`, ```julia const nn = Chain( Base.Fix1(fast_activation, cos), Dense(1 => 32, cos; init_weight=truncated_normal(; std=1.0e-4), init_bias=zeros32), Dense(32 => 32, cos; init_weight=truncated_normal(; std=1.0e-4), init_bias=zeros32), Dense(32 => 2; init_weight=truncated_normal(; std=1.0e-4), init_bias=zeros32), ) ps, st = Lux.setup(Random.default_rng(), nn) ``` ``` ((layer_1 = NamedTuple(), layer_2 = (weight = Float32[-7.936796f-5; -7.364257f-5; 0.00016330472; 6.3057116f-5; -8.211871f-6; -0.00013862252; 0.00015959903; 3.0408835f-5; -0.000105297375; -0.00012277857; -8.5919775f-5; 7.75401f-5; -4.613022f-5; 3.4545205f-6; 0.00015339747; 2.6064039f-5; -0.00013465021; 0.0001826195; -6.248117f-5; 0.00010830289; 0.00010214844; 0.00015743857; 8.535542f-5; 0.00013113461; -6.167884f-5; -0.00012557696; 2.1335121f-5; 0.0001571898; -0.00012408485; 3.871188f-5; -0.00010797259; -0.00011587471;;], bias = Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), layer_3 = (weight = Float32[1.4950157f-5 0.00023673083 -0.00026174507 -5.695939f-5 -2.0938223f-5 0.00011614707 1.2360577f-5 8.429641f-5 -3.455478f-5 -1.8807992f-5 0.0001374548 9.5450625f-5 -8.117573f-5 9.9538345f-5 0.00012858723 -1.7370201f-5 -0.00013494624 1.6174874f-5 1.97591f-5 -2.6970742f-5 -0.00012624837 7.201265f-5 1.7595154f-5 -2.8805585f-5 6.160482f-5 8.530277f-5 0.00011054116 -7.118542f-6 3.698104f-5 3.9176675f-6 5.018404f-5 3.602254f-5; -1.3018777f-5 -8.73679f-5 -6.390673f-5 8.7967026f-5 6.359948f-5 -0.00014688718 -0.00011207426 3.21215f-5 0.00017127748 -0.00013417104 -0.00011215295 0.00015508622 -2.08412f-6 5.4843167f-5 -9.498901f-6 0.00015611794 -3.645631f-5 -0.00017153728 4.7019566f-5 -0.00016717806 -6.294186f-5 -2.4249892f-5 0.00015144802 4.740753f-5 7.587738f-5 9.171603f-5 0.0001388527 0.00011851267 -2.6582584f-5 2.2608992f-5 5.8097823f-5 5.1661505f-6; -9.456473f-6 -0.0001648037 -1.1524755f-5 -5.7739315f-5 -8.9091685f-5 -0.00015889535 -7.277974f-5 3.6573154f-5 -0.00013734064 7.398644f-5 -0.00027732033 -9.312568f-5 7.943617f-5 -2.8512399f-5 -0.00014858105 0.00013128227 -0.00018783618 -0.00015914455 -0.0002267024 -4.084139f-5 8.118561f-5 0.00015385293 -4.5333774f-5 -1.2564536f-6 0.000120291355 0.00010121549 -1.2896347f-5 0.00013242586 3.317706f-6 -0.00024351256 1.9011704f-5 0.00010854192; 4.7781083f-5 -1.9063857f-5 8.151582f-5 0.00017277534 1.3293801f-5 -0.0001630557 8.661659f-5 -3.407816f-5 0.00014170997 7.943988f-5 -0.00012084139 -0.00019894623 5.5695447f-5 0.00012625092 7.0393115f-5 4.76534f-5 0.00036282148 -8.1414284f-5 -1.8611318f-6 -5.2699885f-5 -0.00016924906 5.420088f-5 -0.000121835175 4.6832727f-5 -0.0001221924 2.3367073f-5 5.6275323f-5 5.4620803f-5 0.00021045699 0.00015406877 -0.0001744378 -0.00015024544; 7.4580115f-5 6.9785114f-5 -6.714516f-5 4.1161655f-7 -2.8530972f-5 -4.5051183f-5 -4.1219755f-5 0.00010020237 -0.00010203331 4.4326594f-6 3.761188f-5 -8.8425164f-5 2.718307f-5 2.5635229f-5 -0.00015607491 9.048734f-5 6.919445f-5 -6.729314f-5 -2.9865903f-6 6.860176f-5 6.158385f-5 -1.3717607f-6 5.234855f-5 0.00013657575 -0.00018328674 4.223622f-5 2.8270062f-5 -9.292849f-5 8.4397f-5 -8.580124f-6 -2.3673987f-5 0.0001051048; -7.12364f-5 0.00015983098 7.908807f-5 5.893512f-5 0.00023462997 0.000103693645 9.657947f-5 8.932415f-5 -4.631672f-5 -4.8885802f-5 0.00011728214 -9.9330035f-5 -0.00012069103 -0.0001734228 2.355182f-5 -3.415184f-5 -7.182655f-5 -6.362825f-5 5.3563766f-5 -0.00013036757 0.00020044549 0.00011648791 -0.00014443995 4.35464f-5 -7.753347f-6 -0.0001157376 6.9224894f-5 -0.00016930897 -0.00020152445 -1.0836392f-5 0.00017120512 -1.4395596f-5; 0.000100990845 -8.6833716f-5 -0.00011609131 -2.612782f-5 -4.149975f-6 0.00021802145 7.5687067f-6 -1.1558209f-6 7.2483745f-5 -5.465303f-5 -8.662379f-6 -6.409959f-6 7.022029f-5 0.00015403841 -6.339684f-6 -9.372181f-5 -0.00032630778 8.6751155f-5 -6.6667094f-6 -4.6009474f-5 -3.5893197f-5 -3.880916f-5 -8.59253f-5 -1.9287301f-5 6.4257383f-6 -7.5095995f-5 0.00018511659 2.4880625f-5 0.00010035447 8.1567094f-5 -5.7860554f-5 0.00016781346; 9.82649f-6 1.4792153f-5 -1.9941288f-6 0.00022769089 8.425291f-7 -0.00017833576 -0.00013082915 8.012355f-5 5.354653f-5 -0.0002621904 6.678341f-5 -4.273864f-5 0.00017107172 3.641597f-5 7.272379f-5 -0.00014110502 -0.00014602025 -1.0289417f-5 4.624932f-5 0.00020025784 2.4123097f-5 7.7011486f-5 -4.676713f-5 -7.667115f-5 4.374698f-5 8.4599065f-5 -4.8126978f-5 -9.9782825f-5 -0.00011639356 -4.4723506f-6 5.1392322f-5 -3.0332701f-5; -2.9377577f-5 7.916525f-5 0.00015829095 6.0554685f-5 -4.5555167f-5 8.318181f-5 -0.00014985145 0.00013226026 -4.2185085f-5 -0.00018973416 0.00020522666 3.7278478f-5 3.6340956f-5 -9.246997f-5 3.6233865f-5 -5.3529035f-5 8.982479f-5 2.7368875f-5 0.000114539165 -5.4327127f-5 -7.6402874f-5 -0.00015440834 2.8741239f-5 9.298447f-5 -1.3227844f-5 2.5141127f-5 -0.00013893047 -0.00034913135 0.00011160688 0.00012473289 1.9721838f-5 0.00013002924; 8.832944f-5 7.959618f-5 -4.4581786f-5 -5.727012f-5 0.00021017161 2.3040424f-5 0.000102593986 5.0246086f-5 -0.00010656741 -0.00024881487 2.3423625f-5 -5.9336675f-5 -2.099698f-5 -9.7496624f-5 1.9287954f-6 -0.0001316393 -9.3393f-5 2.3301283f-5 -0.00014160374 -0.00014531125 8.3789026f-5 -4.302865f-6 -1.2589465f-5 -6.4718144f-5 -1.1432865f-5 -2.1344616f-5 7.369199f-5 6.358269f-5 -2.1819773f-5 2.3228275f-5 -6.227244f-6 -0.00022864834; 8.6868495f-5 0.00014558052 2.5887271f-5 7.7617435f-5 1.179677f-5 -9.4206756f-5 6.147397f-6 -1.5925913f-5 7.1712864f-5 1.1808341f-5 -8.625862f-6 -0.00012018939 2.7033579f-5 4.849367f-5 0.00012048228 -0.00031079366 -0.00026895505 9.97321f-6 1.3480943f-5 -0.00014280183 2.4091454f-5 7.241007f-5 -7.813393f-5 5.2400355f-5 -0.00011840859 4.7182297f-5 -7.9966834f-5 0.00014633227 -3.9663973f-6 -0.00012618347 0.0001393391 1.3044078f-5; -0.00012770541 3.799775f-5 2.5284626f-5 -2.937744f-5 1.170753f-5 -3.9474082f-5 5.5164677f-5 -5.294292f-6 0.00012929813 -1.5523729f-5 0.00012781532 -0.00011955024 7.480996f-5 0.0001792476 0.00016853937 0.00012321735 -0.00015790672 0.0001722427 8.2450024f-5 -6.361945f-5 -0.00014806431 -4.360894f-5 -2.5764617f-5 5.8972044f-5 0.00014921205 -7.082219f-6 -7.469288f-5 4.9471737f-7 -0.00013207276 8.757274f-5 0.00010613617 -3.4175664f-5; -4.547526f-5 -7.748395f-5 -0.00016806144 -0.00013935438 0.00014854026 5.258803f-5 -5.7059893f-5 -5.206123f-5 -6.0687187f-5 0.00011735669 -0.0001874987 7.6532604f-5 -2.6521127f-5 2.476936f-5 -0.00011904863 0.000107568296 -0.000118017844 -0.00018054357 0.00013175707 -6.916533f-5 7.2456525f-5 -1.3191645f-6 -0.00038521946 4.813317f-5 -1.4736692f-5 1.3987058f-5 1.3881521f-6 -0.00011228015 2.9862274f-5 4.219824f-6 5.1028648f-5 -8.097454f-5; -3.589522f-5 -9.603833f-5 -3.3071487f-5 2.3457378f-5 -1.2847493f-5 0.000107512584 -2.640358f-5 -0.00022791844 -2.9309556f-5 9.083007f-5 -1.10639685f-5 1.647785f-5 -5.8934387f-5 4.5676374f-5 -8.587353f-6 -0.00011215449 -4.72786f-5 0.00012141605 0.00016855885 -3.925198f-5 -3.0122728f-5 -0.00013491594 -6.2314357f-6 9.014147f-5 -1.705865f-5 6.7199435f-5 1.3581303f-5 -1.1070658f-5 0.00011084768 -0.00017070607 -6.575231f-5 -3.442766f-5; 1.3391874f-5 -8.212596f-5 -3.588902f-5 4.0980107f-5 5.9094218f-5 -0.00012644833 6.75255f-5 3.684539f-5 4.724653f-5 7.693954f-5 -8.000823f-5 -3.414701f-5 -0.00018045046 -0.00013456238 3.1341297f-5 0.00016715923 4.7747584f-5 -5.4263757f-5 -0.00011474392 -6.4578046f-5 4.458281f-5 3.2387823f-5 4.4693497f-5 0.00024022146 1.7655397f-7 0.00013881866 0.000107048705 -6.802194f-5 -1.3766467f-5 1.3145463f-5 1.3998042f-5 -1.5694492f-5; 0.00015743867 0.00015414497 -3.8015176f-5 -2.4108067f-5 6.2637686f-5 1.1878183f-5 7.77206f-5 0.00011013934 -1.2481057f-5 0.00010416297 2.458296f-5 9.681276f-5 -2.896877f-5 4.7478847f-5 -4.2312936f-6 -6.66042f-5 -6.303819f-5 -2.7819713f-5 -8.3029576f-5 0.00010231495 -3.0286303f-5 -0.00010717267 -5.7837344f-5 3.6530088f-5 2.3038205f-5 8.097126f-5 0.00034689507 0.00020583825 -0.000112592556 0.00013693378 -4.251676f-6 6.502289f-5; -4.966114f-5 0.00013897888 0.00010318894 -0.00019471368 6.263541f-5 -0.00015149506 -7.649434f-5 -3.822506f-5 1.8710887f-5 -8.233194f-5 -8.5404296f-5 1.760019f-5 9.849496f-5 -0.00012947046 0.00013546822 -0.000119060955 -1.7070246f-5 2.9467474f-5 -4.570458f-5 0.00016639639 0.00017595054 -5.703227f-5 -7.881526f-6 0.00020069702 -4.5807752f-5 0.000115677496 0.00014545197 -1.8438219f-5 0.00010371841 5.221115f-5 5.4856737f-5 -0.00011054217; -8.761999f-5 -0.00040615085 -0.00016370384 -0.00015833648 -0.000103871214 -0.000109359135 4.9479564f-5 0.0002550029 -2.993675f-6 0.00012291632 0.00016555346 -8.085861f-5 -4.3239866f-5 -4.4671815f-5 6.4877668f-6 -8.399064f-6 -1.9085599f-5 -6.510549f-5 -7.558862f-5 -0.00012819897 -0.00014832217 -0.00019711822 -0.00014878181 0.00017391764 8.902578f-6 5.2665793f-5 -8.9602596f-5 2.3583052f-5 -7.784885f-5 0.00021644418 9.198576f-5 5.187145f-5; -2.0279636f-5 0.00012900053 9.501961f-6 -7.2584706f-5 2.6000675f-5 4.979623f-5 -0.00010225133 4.327667f-5 4.2474534f-5 0.00012328396 0.0001480326 7.72906f-5 -0.00010368979 -3.340996f-5 -6.303696f-5 5.8317022f-5 -0.00018127 9.846218f-6 -1.9912022f-5 6.832871f-5 9.483394f-5 -1.6530872f-5 -5.625166f-5 0.00012781881 -0.00013748105 -0.0001154275 3.19599f-5 -9.985772f-5 -0.00015865256 -1.9288183f-5 0.00018451575 -8.269558f-6; 0.00010246157 -0.00012745793 -5.8699632f-5 4.667096f-5 -0.00017361135 -7.007271f-6 -0.00019125704 -8.665926f-5 2.9797073f-5 1.0633451f-5 7.3212914f-6 -5.606417f-5 -0.00012064606 -5.9839167f-5 -2.9379638f-5 3.612141f-5 4.5457644f-5 -0.00013738414 0.00011335857 -6.5994136f-5 3.1102387f-5 1.0796552f-5 -9.6545096f-5 3.6156354f-5 -1.8567836f-6 -0.0001365158 -2.2672057f-5 -1.2155233f-5 0.00014869995 8.849615f-5 -7.4800235f-5 5.4815453f-5; -9.812964f-6 -6.7544046f-5 -5.5475164f-5 0.000107992964 -8.574461f-5 5.989535f-5 3.395278f-5 0.00016529947 4.58957f-5 3.958388f-5 -0.00015313331 8.569832f-5 -5.879849f-6 4.295306f-5 6.602518f-6 -0.00015162598 2.8028719f-5 3.4649176f-5 5.644887f-5 -0.00013897364 8.201647f-5 4.8751237f-5 -1.42887375f-5 0.00011118326 5.698546f-5 1.3714018f-6 2.5067871f-5 -2.4305331f-5 -0.00010995469 -4.730785f-5 -7.3997835f-5 7.329596f-5; 7.479802f-5 9.009198f-5 -1.582498f-5 -4.31374f-5 0.00012282531 -3.1952084f-5 -7.090515f-5 0.00012982175 7.7616874f-5 0.00013150186 -5.159959f-5 1.41290775f-5 -2.1525486f-6 -6.89365f-5 -8.756183f-5 4.7636247f-5 -0.00010583592 -3.089551f-5 -4.301961f-5 -1.3317541f-5 -8.262396f-5 0.0001233628 0.00012935745 -0.000116236806 -2.6208096f-5 0.00014476175 1.5774561f-6 0.00018920671 7.955148f-5 2.805385f-5 -0.00011374011 -9.007208f-6; 2.0929827f-5 0.00012146997 5.5418805f-5 4.3270742f-5 0.00012297602 4.9863935f-5 -2.8577557f-5 8.723282f-6 6.348473f-5 8.052833f-5 8.418335f-5 -0.00029534457 3.3376913f-5 7.265404f-5 9.9419245f-5 0.00014658761 0.00015614442 -0.00016573531 0.00016078421 6.657062f-5 7.353429f-5 2.3640723f-5 -3.3517917f-5 -0.00023684943 5.0786806f-5 8.1792474f-5 -0.00017925558 -0.000108187705 -9.5980824f-5 -7.08135f-5 -1.6601488f-6 0.00015195954; -4.093456f-5 5.0174534f-5 5.397993f-5 -0.00015914196 1.8501576f-5 2.047878f-5 -5.602667f-5 8.112536f-5 -3.3225577f-5 3.162945f-5 -5.733247f-5 2.405522f-5 -6.3361517f-6 7.341487f-5 -4.582931f-5 7.09245f-5 7.2211486f-5 -4.7283633f-5 -4.731373f-5 -6.7490204f-5 -0.00011442241 3.797742f-5 0.00016731092 -4.3595457f-5 -0.00028058686 -2.1632215f-5 -2.1237494f-5 0.00017202098 8.6820386f-5 -1.6474985f-5 1.9419306f-6 -0.00012640847; -0.00013171151 -0.00019443607 3.9662442f-5 0.000136315 1.1265441f-5 -0.00010103228 -2.3145247f-5 -6.1190617f-6 -4.0175473f-5 -1.7424203f-5 -5.940464f-5 -3.0277484f-5 0.00021988689 -9.799741f-5 0.000111386486 6.5723165f-5 -4.3661097f-5 0.00018785692 7.0347596f-5 -0.00012778852 0.0001938551 9.524316f-5 -0.00016572248 8.034859f-5 0.00016501197 3.5175286f-5 -8.309016f-6 0.0002341072 -2.7306389f-5 0.00016399918 -3.003735f-6 -0.00012518799; -6.924373f-6 6.190074f-5 -3.2250337f-5 -5.7992456f-5 8.466708f-5 -1.3277596f-5 -9.3244314f-5 5.1643954f-5 1.1692727f-5 5.8266767f-5 2.7579157f-5 -8.80305f-5 -5.8420672f-5 5.452205f-5 2.9298995f-5 -0.00013010926 -0.00012952651 0.0001031566 7.321508f-5 -5.1033016f-6 -0.0001037596 2.1814418f-5 0.00016347916 0.00019743382 4.7022904f-6 0.00019068242 -1.294022f-5 3.904697f-5 6.289326f-5 5.392718f-5 -6.3555926f-5 -0.00010543822; 0.00017299916 0.00011655077 4.1409683f-5 -0.00026018024 9.086726f-5 6.4466585f-5 -5.227107f-5 -4.9810806f-6 5.0754858f-5 -4.6128804f-5 0.0001778382 -0.00027820823 2.0657826f-5 -0.00011220161 -3.757358f-5 -0.00015762816 -6.4610955f-5 -4.147493f-6 -3.625072f-5 1.4704098f-5 -0.00019955999 3.1926516f-5 4.644019f-5 -6.805694f-5 6.445848f-5 -8.543875f-5 4.7822803f-5 0.00013661844 -2.2340275f-5 0.0001986795 -0.00019522599 -7.0695554f-5; -3.483146f-5 2.466475f-5 0.00016685853 5.901416f-6 -0.00016996905 9.2970804f-5 -7.831689f-5 4.666848f-5 4.1905678f-5 -0.00017896708 -9.9267825f-5 2.1278924f-5 -4.5819004f-5 -3.190399f-5 0.0002584635 2.2476838f-6 -7.590731f-5 0.00011385473 2.7843742f-5 -7.800999f-5 0.00013918415 -2.1647476f-5 -2.1606898f-5 -0.00016680434 0.00019686285 4.402987f-5 -6.681372f-5 -7.441488f-6 5.0289127f-5 -7.669116f-5 5.860585f-5 9.5456104f-5; -6.559404f-5 -0.00013060244 7.583125f-5 4.5363005f-5 0.00010964736 8.707295f-5 -3.34402f-5 -3.6783487f-5 -2.4163353f-5 -0.00017267893 -0.0001488101 -7.31439f-5 4.9715123f-5 -0.00016379352 8.322829f-5 6.761933f-5 2.9840372f-5 -9.736533f-5 -9.235016f-5 1.693855f-5 0.00015779694 -7.2780225f-5 -1.6330552f-5 0.00013871008 7.692674f-5 -0.00018980511 -1.9240288f-5 0.00029880388 -3.7793998f-7 -0.00012663062 -9.886947f-5 -8.701624f-5; 0.00015363879 -0.00021067062 0.00016489603 -2.753142f-5 -1.2092005f-5 5.13624f-5 5.1311083f-5 0.00013824443 -8.896384f-5 -8.621673f-5 -0.00014458345 0.0001416094 6.5454515f-5 7.209509f-5 -6.641257f-5 -0.00017370422 8.645927f-6 -7.933186f-5 -0.00017470375 3.0552274f-5 2.3044819f-5 1.9475485f-5 -4.0106286f-5 -3.6076297f-5 0.00013721133 -4.9312577f-5 -0.00025021512 -6.4908854f-5 5.7223253f-5 0.0002668476 0.00024140504 0.00022376099; 3.809059f-5 -0.00018468866 2.6164693f-5 0.0001655851 6.5934466f-5 -0.00015996564 -6.396163f-5 -2.9570348f-5 2.5173342f-5 -8.554891f-5 3.5189158f-5 2.5949372f-5 0.00013949841 5.2327785f-5 -6.312468f-6 0.00020136373 0.00028310364 -0.00017677658 -5.8336573f-5 -6.6788365f-5 -0.00020302221 -9.8818f-6 2.4023726f-5 -5.0132483f-5 -0.000179385 -0.00011192273 0.00019104744 4.2862146f-5 -9.322245f-5 3.0287976f-5 -6.851379f-5 -5.00431f-5; -1.5591826f-5 5.0789837f-5 -5.509041f-5 -0.00020979827 -4.522313f-5 -9.3995135f-5 3.686201f-5 0.0001332208 -4.2010004f-5 -5.4289467f-5 -0.00021172519 -8.1277125f-5 8.50264f-5 0.00022232902 6.347964f-5 -0.000103151695 -7.6814336f-5 3.566552f-5 -0.00013502844 0.0002628229 -1.4228726f-5 -2.9511746f-5 -0.00027358526 0.00010237656 0.00025130645 0.00016662659 5.2599687f-5 4.5523f-5 -1.9895459f-5 0.00015163343 -0.00012496466 -9.6571915f-5], bias = Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), layer_4 = (weight = Float32[-6.1188795f-5 -0.00011425061 4.28694f-5 8.822495f-5 0.00018105605 -8.943244f-5 4.2765667f-5 1.349599f-5 -0.00014486739 -6.873888f-6 -5.1785737f-5 0.0002028745 -7.777996f-5 -2.6063517f-5 8.683489f-5 5.950944f-5 5.212992f-5 0.000112782785 7.801872f-5 -7.0091526f-5 -2.2014774f-5 -1.29462605f-5 3.1822947f-5 -1.1362036f-5 0.00013118793 -9.116197f-5 -0.00010286801 -6.7445275f-5 0.00014403173 0.00013638388 -0.00011129156 9.013259f-5; 8.625081f-5 -4.445318f-5 -3.580109f-5 7.375552f-5 3.1194228f-5 -0.0001636009 0.00010799804 -0.00011965589 -0.000116572344 0.00015246005 -0.0002349612 -0.00010699277 2.9763823f-5 2.4058205f-5 0.0001889365 -5.1915078f-5 -2.2433882f-5 -8.758037f-5 -1.0119015f-5 -0.00019676027 0.00013793581 -0.000113158225 -2.2949273f-5 -6.969834f-5 0.00010184694 -3.0066984f-5 6.737199f-5 0.00011478652 -6.164567f-5 3.8741517f-5 0.00011103627 4.5221594f-5], bias = Float32[0.0, 0.0])), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple())) ``` Similar to most deep learning frameworks, Lux defaults to using `Float32`. However, in this case we need Float64 ```julia const params = ComponentArray(f64(ps)) const nn_model = StatefulLuxLayer(nn, nothing, st) ``` ``` StatefulLuxLayer{Val{true}()}( Chain( layer_1 = WrappedFunction(Base.Fix1{typeof(LuxLib.API.fast_activation), typeof(cos)}(LuxLib.API.fast_activation, cos)), layer_2 = Dense(1 => 32, cos), # 64 parameters layer_3 = Dense(32 => 32, cos), # 1_056 parameters layer_4 = Dense(32 => 2), # 66 parameters ), ) # Total: 1_186 parameters, # plus 0 states. ``` Now we define a system of odes which describes motion of point like particle with Newtonian physics, uses $$u\[1] = \chi$$ $$u\[2] = \phi$$ where, $p$, $M$, and $e$ are constants ```julia function ODE_model(u, nn_params, t) χ, ϕ = u p, M, e = ode_model_params # In this example we know that `st` is am empty NamedTuple hence we can safely ignore # it, however, in general, we should use `st` to store the state of the neural network. y = 1 .+ nn_model([first(u)], nn_params) numer = (1 + e * cos(χ))^2 denom = M * (p^(3 / 2)) χ̇ = (numer / denom) * y[1] ϕ̇ = (numer / denom) * y[2] return [χ̇, ϕ̇] end ``` Let us now simulate the neural network model and plot the results. We'll use the untrained neural network parameters to simulate the model. ```julia prob_nn = ODEProblem(ODE_model, u0, tspan, params) soln_nn = Array(solve(prob_nn, RK4(); u0, p=params, saveat=tsteps, dt, adaptive=false)) waveform_nn = first(compute_waveform(dt_data, soln_nn, mass_ratio, ode_model_params)) begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="Time", ylabel="Waveform") l1 = lines!(ax, tsteps, waveform; linewidth=2, alpha=0.75) s1 = scatter!( ax, tsteps, waveform; marker=:circle, markersize=12, alpha=0.5, strokewidth=2 ) l2 = lines!(ax, tsteps, waveform_nn; linewidth=2, alpha=0.75) s2 = scatter!( ax, tsteps, waveform_nn; marker=:circle, markersize=12, alpha=0.5, strokewidth=2 ) axislegend( ax, [[l1, s1], [l2, s2]], ["Waveform Data", "Waveform Neural Net (Untrained)"]; position=:lb, ) fig end ``` ## Setting Up for Training the Neural Network {#Setting-Up-for-Training-the-Neural-Network} Next, we define the objective (loss) function to be minimized when training the neural differential equations. ```julia const mseloss = MSELoss() function loss(θ) pred = Array(solve(prob_nn, RK4(); u0, p=θ, saveat=tsteps, dt, adaptive=false)) pred_waveform = first(compute_waveform(dt_data, pred, mass_ratio, ode_model_params)) return mseloss(pred_waveform, waveform) end ``` Warmup the loss function ```julia loss(params) ``` ``` 0.0007333700155341713 ``` Now let us define a callback function to store the loss over time ```julia const losses = Float64[] function callback(θ, l) push!(losses, l) @printf "Training \t Iteration: %5d \t Loss: %.10f\n" θ.iter l return false end ``` ## Training the Neural Network {#Training-the-Neural-Network} Training uses the BFGS optimizers. This seems to give good results because the Newtonian model seems to give a very good initial guess ```julia adtype = Optimization.AutoZygote() optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype) optprob = Optimization.OptimizationProblem(optf, params) res = Optimization.solve( optprob, BFGS(; initial_stepnorm=0.01, linesearch=LineSearches.BackTracking()); callback, maxiters=1000, ) ``` ``` retcode: Success u: ComponentVector{Float64}(layer_1 = Float64[], layer_2 = (weight = [-7.936795736893456e-5; -7.364257180576249e-5; 0.00016330472135435358; 6.305711576707124e-5; -8.21187131804941e-6; -0.0001386225194437058; 0.00015959903248575135; 3.0408835300455e-5; -0.00010529737482998399; -0.00012277856876607247; -8.591977530151319e-5; 7.754009857302936e-5; -4.613021883413124e-5; 3.454520538066933e-6; 0.0001533974718765434; 2.60640390479281e-5; -0.0001346502103841956; 0.00018261949298873024; -6.24811727902364e-5; 0.00010830289102148378; 0.0001021484422381269; 0.0001574385678395351; 8.53554229250815e-5; 0.00013113461318428087; -6.167883839221397e-5; -0.0001255769602719877; 2.1335121346030356e-5; 0.0001571898028487354; -0.00012408485054041322; 3.8711881643333754e-5; -0.0001079725916497156; -0.0001158747109001618;;], bias = [-2.735476363621604e-17, -2.457431976093982e-16, 1.7530205584028062e-16, 9.024817478322839e-17, -4.816481364711468e-18, -4.3763393718579855e-17, 1.224720462065248e-16, 2.5609352476166116e-17, -1.4946836390571246e-16, -6.7483427234965855e-18, -8.936918554638135e-17, 5.988280413757502e-17, -2.1870956854437202e-17, 4.545591085371465e-18, 3.0835889520115894e-16, -4.94886411529186e-18, -9.50468070773812e-17, 3.4628098429228793e-16, -5.497320368253868e-17, 1.2707527856850822e-16, -2.7449861025882532e-17, 3.0136370111216595e-17, 1.4217734648022492e-16, -2.745376561252617e-17, -4.536404813227628e-17, -1.9244066821807871e-16, 2.886480992435844e-17, -2.638044620437127e-18, 3.441421892485586e-17, 6.808304611389256e-17, -2.870432870682115e-17, -9.197945482228104e-17]), layer_3 = (weight = [1.4952754513434836e-5 0.00023673343135197 -0.0002617424684515002 -5.6956792418710546e-5 -2.0935625384656755e-5 0.00011614966975211209 1.236317480541739e-5 8.429900522497218e-5 -3.455218351992045e-5 -1.880539416125116e-5 0.00013745739893674966 9.545322287049377e-5 -8.1173129315352e-5 9.954094316963849e-5 0.0001285898255944499 -1.7367603091036537e-5 -0.00013494364206527226 1.6177471893366896e-5 1.976169774461952e-5 -2.6968143670459017e-5 -0.0001262457731583729 7.201524609152724e-5 1.759775155260353e-5 -2.8802987385650838e-5 6.160741800595499e-5 8.53053647504501e-5 0.00011054375541317051 -7.115943885960933e-6 3.69836383968442e-5 3.920265457787688e-6 5.018663837726912e-5 3.602513648256262e-5; -1.3017490230250982e-5 -8.736661481873037e-5 -6.390544228691475e-5 8.79683130982627e-5 6.360076725048336e-5 -0.00014688589078343955 -0.00011207297290878187 3.2122788186262414e-5 0.0001712787699283366 -0.00013416975786879038 -0.00011215166239037166 0.0001550875086899202 -2.0828330185572663e-6 5.4844454214054605e-5 -9.497613857665248e-6 0.0001561192249276984 -3.6455022266297336e-5 -0.00017153599116921577 4.702085346014398e-5 -0.00016717677578765138 -6.294057027173344e-5 -2.42486047433249e-5 0.00015144931160409952 4.740881843406307e-5 7.587866923802697e-5 9.171731574258426e-5 0.0001388539886896808 0.00011851395459177188 -2.6581296763311592e-5 2.2610279418023125e-5 5.809911010817457e-5 5.167437509385552e-6; -9.461305140512197e-6 -0.0001648085317501059 -1.152958704178199e-5 -5.774414695764245e-5 -8.909651752759671e-5 -0.00015890017768097794 -7.278456998019179e-5 3.6568321760624444e-5 -0.00013734547339110735 7.398160867895527e-5 -0.0002773251626567145 -9.313051215570386e-5 7.943133731174637e-5 -2.851723094494525e-5 -0.0001485858820305113 0.00013127743930823298 -0.0001878410121352404 -0.0001591493792292415 -0.00022670722775359006 -4.08462228878494e-5 8.118077583647133e-5 0.0001538481001117379 -4.5338606483645245e-5 -1.261285788960928e-6 0.00012028652325540468 0.00010121065768915868 -1.2901178721167621e-5 0.00013242103067032243 3.3128737723797933e-6 -0.00024351739649859002 1.9006872193203214e-5 0.00010853708862130799; 4.77831283815372e-5 -1.906181090179034e-5 8.151786575308379e-5 0.00017277738631220498 1.3295846899398496e-5 -0.00016305364873278076 8.661863218276222e-5 -3.407611323653556e-5 0.0001417120158903436 7.944192590640079e-5 -0.00012083934301061821 -0.0001989441868617962 5.5697493068707085e-5 0.00012625296337047818 7.0395160279987e-5 4.7655446239348455e-5 0.0003628235219976022 -8.141223837174609e-5 -1.8590860536139088e-6 -5.269783899438191e-5 -0.000169247016613544 5.420292406324734e-5 -0.0001218331296813514 4.6834772790139666e-5 -0.00012219035737233262 2.336911826842571e-5 5.627736870066459e-5 5.462284868099237e-5 0.00021045903237518985 0.00015407081197464965 -0.00017443574937077793 -0.0001502433940304324; 7.458085922293208e-5 6.978585763603314e-5 -6.714441958580812e-5 4.1236045275526e-7 -2.853022858731454e-5 -4.505043950745371e-5 -4.1219011177296427e-5 0.00010020311485744987 -0.00010203256358623777 4.433403251038215e-6 3.761262444198881e-5 -8.842442009596304e-5 2.7183814313630022e-5 2.5635972565524244e-5 -0.00015607416396493662 9.048808162800948e-5 6.919519539691066e-5 -6.729239801176701e-5 -2.9858463902752165e-6 6.860250316561846e-5 6.158459466417996e-5 -1.3710168330964146e-6 5.2349294424073836e-5 0.00013657649280970495 -0.000183285998059419 4.223696480506647e-5 2.827080600140136e-5 -9.29277430940293e-5 8.439774130705368e-5 -8.579380005230596e-6 -2.3673243117502945e-5 0.00010510554595851942; -7.123600242119846e-5 0.00015983137111253486 7.90884687572635e-5 5.893551408154357e-5 0.00023463036906979493 0.00010369404028243588 9.6579863585803e-5 8.932454786337419e-5 -4.6316324167844574e-5 -4.8885406593752336e-5 0.0001172825386871526 -9.93296389227789e-5 -0.00012069063539309935 -0.00017342240333017444 2.3552215699610826e-5 -3.415144326789796e-5 -7.18261553432864e-5 -6.362785368618537e-5 5.356416206415211e-5 -0.0001303671715308044 0.00020044588293022792 0.0001164883024299437 -0.00014443955749665033 4.354679460594658e-5 -7.752951185818022e-6 -0.00011573720710511074 6.922528975484405e-5 -0.0001693085768951183 -0.00020152405793704773 -1.0835996343654474e-5 0.00017120551185290626 -1.4395200176648364e-5; 0.00010099161221784555 -8.683294805326461e-5 -0.00011609054132226419 -2.61270525321683e-5 -4.149207651273493e-6 0.00021802221812446297 7.569474238480909e-6 -1.1550533223567881e-6 7.248451214888026e-5 -5.4652260640647456e-5 -8.661611044439606e-6 -6.409191481574447e-6 7.022105632255906e-5 0.00015403917854325084 -6.338916644962249e-6 -9.372103891894995e-5 -0.0003263070139461317 8.675192288261175e-5 -6.665941835884416e-6 -4.600870675734364e-5 -3.5892429931332304e-5 -3.880839363649936e-5 -8.592453019321904e-5 -1.9286533696356014e-5 6.426505880305392e-6 -7.509522771091425e-5 0.00018511735465050837 2.4881392111056675e-5 0.00010035523513703392 8.156786129017498e-5 -5.785978651932559e-5 0.00016781422522523277; 9.82616870018823e-6 1.4791831510906778e-5 -1.9944500378661526e-6 0.00022769057135732992 8.42207832506418e-7 -0.00017833607967773148 -0.00013082947521766275 8.01232297293108e-5 5.3546208384747004e-5 -0.00026219073448160394 6.67830888617988e-5 -4.273896161762831e-5 0.00017107139751019088 3.641565012216672e-5 7.272347169632442e-5 -0.00014110533926000703 -0.00014602056877052117 -1.0289738618349794e-5 4.6248997698371817e-5 0.00020025751904754194 2.4122776032842025e-5 7.701116441453239e-5 -4.676745198380954e-5 -7.667147406905605e-5 4.374665763165945e-5 8.459874401737273e-5 -4.8127299358070676e-5 -9.978314597688254e-5 -0.00011639388072367492 -4.472671896705839e-6 5.139200106200055e-5 -3.033302278495616e-5; -2.93767244349497e-5 7.916610220290907e-5 0.00015829179928714007 6.0555537764050465e-5 -4.5554314149009437e-5 8.318266363657815e-5 -0.0001498505936089754 0.00013226111504791618 -4.218423246329621e-5 -0.0001897333041617469 0.0002052275175608894 3.7279330988714315e-5 3.6341809298212636e-5 -9.246911768259143e-5 3.62347181160623e-5 -5.3528181617545506e-5 8.982564204215836e-5 2.7369727738726545e-5 0.00011454001823749623 -5.432627413229074e-5 -7.640202066086185e-5 -0.0001544074822069926 2.8742091579080043e-5 9.298532122356675e-5 -1.3226990724613437e-5 2.5141980493388976e-5 -0.00013892961406072275 -0.00034913050158338974 0.00011160773455940095 0.0001247337439943321 1.9722690829257604e-5 0.0001300300881646756; 8.832661381721205e-5 7.959335465248302e-5 -4.4584613593428256e-5 -5.7272948712031775e-5 0.00021016878139025594 2.303759623293274e-5 0.00010259115782366689 5.024325849635667e-5 -0.00010657023546254108 -0.00024881770247328497 2.3420797273598567e-5 -5.933950257339833e-5 -2.0999807989846247e-5 -9.749945190005878e-5 1.9259677484319772e-6 -0.00013164213038717278 -9.339582635755776e-5 2.3298455684318952e-5 -0.0001416065688917092 -0.0001453140767497548 8.3786197987237e-5 -4.305692587764678e-6 -1.2592292538005534e-5 -6.472097180984195e-5 -1.1435692673237637e-5 -2.134744414969227e-5 7.368916243112745e-5 6.357985961051211e-5 -2.1822600561686863e-5 2.3225446906599295e-5 -6.2300715088800506e-6 -0.00022865116358469805; 8.686756957137605e-5 0.00014557959259070083 2.58863463571973e-5 7.761650960707615e-5 1.1795845075872923e-5 -9.420768144957685e-5 6.146471896840865e-6 -1.5926837640684577e-5 7.171193903593138e-5 1.1807415667473383e-5 -8.626787511849248e-6 -0.00012019031526472782 2.7032653470421093e-5 4.849274666445721e-5 0.00012048135142356324 -0.00031079458965872 -0.0002689559707320108 9.97228455196595e-6 1.3480018260786722e-5 -0.00014280275461681376 2.409052905965834e-5 7.240914312239938e-5 -7.81348582492056e-5 5.239943022422468e-5 -0.00011840951736269945 4.7181371719827195e-5 -7.996775928371182e-5 0.000146331344531404 -3.967322463384688e-6 -0.00012618439466644576 0.00013933817428584953 1.304315248523468e-5; -0.0001277029060837743 3.8000254662715105e-5 2.5287130356523936e-5 -2.9374934583660702e-5 1.1710034286906356e-5 -3.9471577238062154e-5 5.5167181799412535e-5 -5.2917872743742785e-6 0.00012930063699977782 -1.5521223885938674e-5 0.00012781782594184068 -0.00011954773630340797 7.481246744656103e-5 0.00017925010057307119 0.00016854187002544524 0.000123219857287139 -0.00015790421288321057 0.00017224520148381963 8.2452528531672e-5 -6.361694808210922e-5 -0.00014806180673977 -4.360643482860977e-5 -2.576211240002271e-5 5.8974548537747934e-5 0.0001492145546400078 -7.0797142506351315e-6 -7.469037349427814e-5 4.972219904514354e-7 -0.00013207025598594956 8.757524652547354e-5 0.00010613867367292085 -3.417315941979058e-5; -4.547951530770584e-5 -7.748820635828498e-5 -0.00016806569320252304 -0.00013935863526643767 0.00014853600450219468 5.258377504431396e-5 -5.706414950182967e-5 -5.206548658538613e-5 -6.069144356518501e-5 0.00011735243614891653 -0.00018750296285946887 7.652834764809123e-5 -2.6525383788289724e-5 2.476510425690431e-5 -0.00011905288548078963 0.00010756403943883218 -0.0001180221005655958 -0.000180547831320289 0.00013175281500803717 -6.916958946606503e-5 7.245226888114045e-5 -1.3234208894671696e-6 -0.00038522372042865534 4.812891356193865e-5 -1.4740948160072531e-5 1.3982801732322845e-5 1.38389574981546e-6 -0.00011228440408232101 2.9858018109348264e-5 4.215567599706887e-6 5.1024391808431624e-5 -8.097879609197504e-5; -3.5896714862682006e-5 -9.603982779524712e-5 -3.307298119622599e-5 2.3455883464994556e-5 -1.2848987708350043e-5 0.00010751108912657265 -2.640507535986046e-5 -0.00022791993559665694 -2.9311050994216067e-5 9.082857705791528e-5 -1.1065463141201546e-5 1.6476355639971573e-5 -5.893588185288596e-5 4.5674879047072136e-5 -8.588847222626252e-6 -0.00011215598659917436 -4.728009606921847e-5 0.00012141455646045516 0.00016855735576384203 -3.925347427000407e-5 -3.0124222931113943e-5 -0.0001349174375631887 -6.2329303459390685e-6 9.013997315335239e-5 -1.7060144628412453e-5 6.719794019729913e-5 1.3579808737193559e-5 -1.1072152474724622e-5 0.00011084618437009925 -0.00017070756747714858 -6.575380702131546e-5 -3.4429154211897726e-5; 1.3392807584607892e-5 -8.212502628535432e-5 -3.588808734290421e-5 4.098104060889798e-5 5.90951519078272e-5 -0.0001264473952070706 6.752643336570739e-5 3.684632489952927e-5 4.724746274209007e-5 7.69404747805675e-5 -8.000729333812832e-5 -3.414607757091009e-5 -0.0001804495235157043 -0.00013456144127499552 3.1342230311417736e-5 0.00016716016640916208 4.774851792520986e-5 -5.4262823443574323e-5 -0.00011474298419867065 -6.457711181806994e-5 4.458374193549533e-5 3.2388756760917e-5 4.469443107868086e-5 0.00024022239328391087 1.7748771685439163e-7 0.00013881959845806686 0.0001070496388527726 -6.802100445580728e-5 -1.376553366950254e-5 1.3146396366068615e-5 1.399897579889619e-5 -1.5693557845369846e-5; 0.00015744377161703115 0.0001541500765723288 -3.801007404599878e-5 -2.410296509529965e-5 6.264278782859593e-5 1.1883284618135066e-5 7.772570244360374e-5 0.0001101444446374849 -1.247595481452117e-5 0.00010416807491657775 2.458806158195581e-5 9.681785887899901e-5 -2.8963668579451825e-5 4.7483948595842046e-5 -4.22619167347573e-6 -6.659910139765545e-5 -6.303308908735456e-5 -2.7814611154294268e-5 -8.302447384827326e-5 0.00010232005444215686 -3.0281200803231333e-5 -0.00010716757095560659 -5.783224186067587e-5 3.653518944139693e-5 2.304330665887327e-5 8.097636019903572e-5 0.00034690017638590277 0.0002058433543550574 -0.0001125874536102691 0.0001369388861474758 -4.246573904470128e-6 6.502799225365538e-5; -4.965946285831752e-5 0.00013898055763836056 0.00010319061689929848 -0.00019471200070422955 6.263708488238376e-5 -0.00015149338000068226 -7.64926633945306e-5 -3.822338189077281e-5 1.8712563504019837e-5 -8.23302660517909e-5 -8.540261918990941e-5 1.760186674634466e-5 9.849663470792803e-5 -0.0001294687798161492 0.00013546989353759194 -0.00011905927826030813 -1.706856968834254e-5 2.9469150065918495e-5 -4.570290260910868e-5 0.00016639806850268196 0.0001759522157573583 -5.703059363608335e-5 -7.879849219381048e-6 0.0002006986934776823 -4.580607568807775e-5 0.00011567917241169545 0.000145453642433652 -1.84365424882722e-5 0.00010372008833488874 5.221282768449998e-5 5.4858413166736036e-5 -0.00011054049249818017; -8.762408572194935e-5 -0.00040615494538044503 -0.00016370793048235517 -0.000158340572965724 -0.00010387530812707461 -0.00010936322826137755 4.947547029178501e-5 0.0002549988083733004 -2.9977687741096035e-6 0.0001229122234467363 0.0001655493641696715 -8.086270565811419e-5 -4.324396024870145e-5 -4.467590872494243e-5 6.483673019375445e-6 -8.40315817598138e-6 -1.9089692727454395e-5 -6.510958075931271e-5 -7.55927150062487e-5 -0.00012820306176264824 -0.00014832626504768842 -0.00019712231471612294 -0.00014878590184209573 0.00017391354310847563 8.898484225457708e-6 5.266169944251443e-5 -8.960668958437312e-5 2.357895834893729e-5 -7.785294030737626e-5 0.00021644008927566453 9.198166924350837e-5 5.186735768391927e-5; -2.0279567020153778e-5 0.00012900060051386234 9.502030711771777e-6 -7.258463642294528e-5 2.6000743905641126e-5 4.9796297621540276e-5 -0.00010225126220113656 4.3276737646991365e-5 4.247460333779428e-5 0.0001232840261355505 0.00014803267209124736 7.729066806412289e-5 -0.00010368971902146009 -3.34098910979461e-5 -6.303689201097057e-5 5.8317091547977164e-5 -0.00018126993722482625 9.846287189318506e-6 -1.991195289967591e-5 6.832877648328404e-5 9.483400595406979e-5 -1.6530802663439518e-5 -5.625159223647805e-5 0.00012781888313391208 -0.00013748098330772487 -0.00011542743208785577 3.195996783240697e-5 -9.985765404501054e-5 -0.00015865249201389993 -1.9288114112825172e-5 0.00018451582241603147 -8.269488484913331e-6; 0.00010245860850829184 -0.00012746088812435577 -5.8702593905184175e-5 4.6667998886381066e-5 -0.0001736143099873432 -7.010232893090185e-6 -0.000191260006048982 -8.66622218194514e-5 2.979411094258553e-5 1.063048914204516e-5 7.318329535898576e-6 -5.606713291798263e-5 -0.00012064902019063859 -5.984212892089161e-5 -2.9382600204704318e-5 3.611844969837868e-5 4.545468202252633e-5 -0.00013738709874756342 0.00011335560650023883 -6.599709793534054e-5 3.109942501452702e-5 1.0793589736418155e-5 -9.654805815387871e-5 3.6153392484836716e-5 -1.8597454112630355e-6 -0.0001365187568620732 -2.267501857149028e-5 -1.2158194821826902e-5 0.00014869699172834363 8.849318639532683e-5 -7.480319671174442e-5 5.481249119895749e-5; -9.812258968917418e-6 -6.75433410206768e-5 -5.5474459105497974e-5 0.00010799366934309233 -8.574390160999652e-5 5.989605339793586e-5 3.3953484643389925e-5 0.00016530017380221332 4.589640518036025e-5 3.958458470963134e-5 -0.00015313260594764354 8.569902199171886e-5 -5.8791441288785664e-6 4.2953764176606956e-5 6.603223181556653e-6 -0.00015162527491254204 2.8029423600785727e-5 3.464988109855683e-5 5.6449574157274196e-5 -0.00013897293759425837 8.201717643617022e-5 4.8751942057534745e-5 -1.428803246296802e-5 0.0001111839657541379 5.698616511738382e-5 1.3721068343543685e-6 2.5068576332093873e-5 -2.430462594118465e-5 -0.00010995398845905252 -4.730714575186796e-5 -7.399712997637415e-5 7.329666474918137e-5; 7.480003654148879e-5 9.009399030713872e-5 -1.5822966874637166e-5 -4.3135385626157315e-5 0.0001228273230803471 -3.1950070537798344e-5 -7.090313926217022e-5 0.00012982376750682352 7.76188880403726e-5 0.00013150387343142055 -5.1597575351992565e-5 1.4131091061641958e-5 -2.1505350692317457e-6 -6.893448341048393e-5 -8.755981440533632e-5 4.763826103539093e-5 -0.00010583390582230267 -3.089349599084277e-5 -4.301759514834125e-5 -1.3315527271122196e-5 -8.262194939067408e-5 0.00012336481262120586 0.00012935945954762785 -0.00011623479264432825 -2.620608217017577e-5 0.00014476376776240593 1.5794697087047415e-6 0.00018920872201799588 7.955349241040501e-5 2.805586330890954e-5 -0.00011373809872072093 -9.005194501304239e-6; 2.093219138957672e-5 0.00012147233776674786 5.5421169079322697e-5 4.3273106555420084e-5 0.00012297838823331183 4.986629922101323e-5 -2.8575192275042323e-5 8.725646537347364e-6 6.348709560158405e-5 8.053069161875539e-5 8.418571799454309e-5 -0.00029534220099988257 3.337927758154891e-5 7.265640290613878e-5 9.942160961841882e-5 0.0001465899766368939 0.00015614678689204284 -0.00016573294867371319 0.00016078657775112727 6.657298202086866e-5 7.353665553425232e-5 2.364308768633838e-5 -3.351555294314363e-5 -0.00023684706700108742 5.07891707708717e-5 8.179483832251594e-5 -0.00017925322042385483 -0.00010818534080712497 -9.597845944184467e-5 -7.081113305714435e-5 -1.65778423918916e-6 0.0001519619034549271; -4.093550870827546e-5 5.0173585065135854e-5 5.3978980018957686e-5 -0.000159142906171542 1.8500626707758147e-5 2.0477829990640934e-5 -5.6027620133743826e-5 8.112441372424981e-5 -3.322652611466124e-5 3.162850060682867e-5 -5.733341805687289e-5 2.4054270567217205e-5 -6.337101070745122e-6 7.341392136943184e-5 -4.583026003608095e-5 7.092355025299502e-5 7.221053618771244e-5 -4.7284582081979706e-5 -4.731467908065988e-5 -6.749115328936817e-5 -0.00011442335727558492 3.797647167296573e-5 0.00016730997248153086 -4.359640644331273e-5 -0.0002805878064588573 -2.163316440935101e-5 -2.1238443708786302e-5 0.00017202003453924743 8.68194367161127e-5 -1.6475934729602355e-5 1.9409812342266733e-6 -0.00012640942060322425; -0.0001317088833274831 -0.00019443344417591958 3.966506783401259e-5 0.00013631761912624095 1.1268066646194402e-5 -0.00010102965525217978 -2.3142620830165667e-5 -6.116435724793651e-6 -4.017284727521173e-5 -1.7421577194880638e-5 -5.940201363042545e-5 -3.02748582725689e-5 0.00021988951566531226 -9.799478057171136e-5 0.00011138911149830492 6.572579086427602e-5 -4.365847116784728e-5 0.00018785954243868616 7.035022217682314e-5 -0.0001277858907087332 0.00019385772548051672 9.524578426864758e-5 -0.0001657198524298666 8.035121509240487e-5 0.00016501459258510493 3.517791201377886e-5 -8.306390282045598e-6 0.0002341098218494358 -2.7303762817581402e-5 0.000164001808389043 -3.0011089620213504e-6 -0.00012518536159401893; -6.9227388180161426e-6 6.190237461263267e-5 -3.224870245183139e-5 -5.7990822212071796e-5 8.46687150201784e-5 -1.3275961761312486e-5 -9.324268041971326e-5 5.164558813115776e-5 1.169436079501989e-5 5.8268401220202405e-5 2.7580791241720666e-5 -8.80288674365737e-5 -5.841903778351714e-5 5.452368411491867e-5 2.930062935186e-5 -0.0001301076220683733 -0.00012952487607114746 0.00010315823299017127 7.321671425737477e-5 -5.101667581100832e-6 -0.00010375796712329925 2.181605184298324e-5 0.0001634807892274097 0.00019743545696976846 4.703924490546737e-6 0.00019068405224382276 -1.2938585611174491e-5 3.904860266040863e-5 6.28948952667357e-5 5.392881273431866e-5 -6.355429225229105e-5 -0.0001054365887525539; 0.0001729974170390726 0.0001165490244690533 4.140793645063882e-5 -0.0002601819873146517 9.086551254737995e-5 6.446483858752104e-5 -5.227281543767612e-5 -4.982827149732818e-6 5.075311092666855e-5 -4.6130550245216226e-5 0.0001778364527214534 -0.00027820997674127826 2.065607951378441e-5 -0.00011220335763057824 -3.7575326487344604e-5 -0.00015762991132790315 -6.461270130814076e-5 -4.1492397852459655e-6 -3.6252466443627136e-5 1.4702351663506885e-5 -0.00019956173527163918 3.192476980046728e-5 4.643844259560394e-5 -6.805868942166755e-5 6.445673317072999e-5 -8.544049377662286e-5 4.782105644123523e-5 0.00013661669764559136 -2.234202146639295e-5 0.00019867774936726274 -0.000195227738359174 -7.069730089853812e-5; -3.4830486532248283e-5 2.4665724646818444e-5 0.00016685950600430304 5.902390109166353e-6 -0.00016996807329933302 9.29717780824768e-5 -7.831591296354337e-5 4.666945374692616e-5 4.1906651910445326e-5 -0.0001789661045616462 -9.926685101463838e-5 2.1279897788824875e-5 -4.581803017872598e-5 -3.1903015900116e-5 0.0002584644672030586 2.2486579443462875e-6 -7.590633408045189e-5 0.00011385570456394646 2.7844716038724515e-5 -7.800901307137272e-5 0.00013918512473237222 -2.1646502267057134e-5 -2.1605924251437233e-5 -0.00016680336643121898 0.00019686382352729063 4.403084590803184e-5 -6.68127476667908e-5 -7.440513803256935e-6 5.029010122297627e-5 -7.669018754613777e-5 5.860682272823011e-5 9.545707785645858e-5; -6.559584044643842e-5 -0.00013060423707847698 7.582945062987597e-5 4.536120349748689e-5 0.00010964555503542005 8.707114711379289e-5 -3.3442001317404e-5 -3.678528928922037e-5 -2.4165155359328524e-5 -0.0001726807361634134 -0.00014881189900242827 -7.314570148500794e-5 4.971332118238844e-5 -0.00016379532217306617 8.322648745471519e-5 6.761752838010406e-5 2.983857002201083e-5 -9.736713155196498e-5 -9.23519594876778e-5 1.6936747393088064e-5 0.00015779513586985954 -7.278202729556728e-5 -1.6332354157018625e-5 0.00013870827633164358 7.692494063612693e-5 -0.00018980691429950844 -1.9242089642468324e-5 0.00029880207393527005 -3.797419693387087e-7 -0.00012663242278413358 -9.887126844168493e-5 -8.701804227220574e-5; 0.00015364027083839685 -0.0002106691387948816 0.00016489751604372978 -2.752993506740273e-5 -1.2090520357128744e-5 5.1363884523620394e-5 5.1312567194564556e-5 0.000138245916133632 -8.896235496516176e-5 -8.621524440835259e-5 -0.00014458196503121994 0.0001416108864591218 6.545599924854594e-5 7.209657658811424e-5 -6.64110845887425e-5 -0.0001737027339172762 8.64741130632404e-6 -7.933037679189796e-5 -0.00017470226131856864 3.0553758080230484e-5 2.304630314162661e-5 1.9476969356638986e-5 -4.010480162699361e-5 -3.607481241354851e-5 0.0001372128174639079 -4.931109262409684e-5 -0.0002502136337339158 -6.490736970457742e-5 5.7224737503444335e-5 0.0002668490926966786 0.00024140652411139708 0.00022376247241615363; 3.808954788890384e-5 -0.00018468970229097953 2.616364945478223e-5 0.00016558405041048875 6.593342259346012e-5 -0.00015996668225799524 -6.396266980974559e-5 -2.95713918598462e-5 2.5172298410854982e-5 -8.554995583782885e-5 3.518811427109629e-5 2.5948328403119802e-5 0.00013949737040145835 5.232674157066703e-5 -6.313511402743571e-6 0.00020136268431140568 0.00028310259658663657 -0.00017677762225449373 -5.833761606429137e-5 -6.678940844669101e-5 -0.00020302325283357082 -9.882843368734437e-6 2.4022682555897475e-5 -5.0133526382901726e-5 -0.00017938603850727335 -0.00011192377423814455 0.00019104639504855948 4.286110270239762e-5 -9.3223493603924e-5 3.0286932815741938e-5 -6.85148363112761e-5 -5.004414488159158e-5; -1.5591922886242257e-5 5.078973980488615e-5 -5.50905083517222e-5 -0.00020979836200495824 -4.522322570920579e-5 -9.399523197572275e-5 3.68619125649189e-5 0.00013322070798239087 -4.201010098114232e-5 -5.4289563661595624e-5 -0.00021172528296375278 -8.127722186401125e-5 8.502629997309647e-5 0.0002223289243258384 6.347954102892848e-5 -0.00010315179180252952 -7.681443322581771e-5 3.5665421353033564e-5 -0.0001350285406377317 0.0002628228007662892 -1.4228823158928163e-5 -2.9511842979272347e-5 -0.0002735853581842186 0.00010237646118838738 0.00025130635684689857 0.0001666264920342567 5.259959060566844e-5 4.552290328064423e-5 -1.9895555408593085e-5 0.00015163333363233484 -0.00012496475620792793 -9.657201236478599e-5], bias = [2.597930513338009e-9, 1.287050855795338e-9, -4.8321598486759125e-9, 2.0456964359003464e-9, 7.438985816760811e-10, 3.9564985299531496e-10, 7.675602024995543e-10, -3.212859657463946e-10, 8.530130829899215e-10, -2.827678964850806e-9, -9.251194873934326e-10, 2.5046211787737096e-9, -4.2563501199888825e-9, -1.4946549361468478e-9, 9.337489655949028e-10, 5.101914072127963e-9, 1.6763425931188783e-9, -4.093748915222013e-9, 6.935358978168804e-11, -2.961841701831603e-9, 7.050120325905316e-10, 2.0135625538050314e-9, 2.364529498298558e-9, -9.49343008619403e-10, 2.6259840940139492e-9, 1.6340650642379526e-9, -1.746584830254639e-9, 9.741203938207127e-10, -1.8019860384320006e-9, 1.48455135290407e-9, -1.0433717960047142e-9, -9.68728366508205e-11]), layer_4 = (weight = [-0.0007679744771765036 -0.0008210363923052917 -0.0006639159694635609 -0.0006185607862004901 -0.0005257297577583303 -0.0007962182549491301 -0.0006640201393879177 -0.0006932898257106355 -0.0008516531895721771 -0.0007136595466233263 -0.0007585715381923974 -0.0005039112113848081 -0.0007845654109273597 -0.0007328492905926212 -0.000619950913802713 -0.0006472758829870859 -0.0006546558449978033 -0.0005940027197803638 -0.0006287670993205304 -0.000776877172233921 -0.0007288005818182619 -0.0007197319997176535 -0.0006749627631483912 -0.00071814783648625 -0.0005755977558348991 -0.0007979477359300899 -0.0008096537681848693 -0.0007742310734446804 -0.0005627540293688892 -0.0005704018944898139 -0.0008180773587568896 -0.000616653226434019; 0.00032480690039159073 0.00019410294556348751 0.0002027548947820168 0.0003123116261135261 0.0002697503607305089 7.495523901182081e-5 0.0003465541756247066 0.0001189002480417379 0.00012198378731759597 0.0003910161350353019 3.594926852173615e-6 0.0001315633308738516 0.000268319837472898 0.0002626143265672063 0.0004274926247950857 0.000186640891778824 0.00021612223654404137 0.0001509756643720869 0.00022843712119161754 4.179580589026711e-5 0.00037649194776069134 0.0001253978852789259 0.00021560682775684554 0.0001688577899740924 0.00034040303376927327 0.0002084891349903316 0.00030592810710072934 0.0003533426487580017 0.00017691044902731274 0.0002772976392541454 0.0003495924008662978 0.00028377773008610255], bias = [-0.0007067858177529196, 0.00023855613653692674])) ``` ## Visualizing the Results {#Visualizing-the-Results} Let us now plot the loss over time ```julia begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="Iteration", ylabel="Loss") lines!(ax, losses; linewidth=4, alpha=0.75) scatter!(ax, 1:length(losses), losses; marker=:circle, markersize=12, strokewidth=2) fig end ``` Finally let us visualize the results ```julia prob_nn = ODEProblem(ODE_model, u0, tspan, res.u) soln_nn = Array(solve(prob_nn, RK4(); u0, p=res.u, saveat=tsteps, dt, adaptive=false)) waveform_nn_trained = first( compute_waveform(dt_data, soln_nn, mass_ratio, ode_model_params) ) begin fig = Figure() ax = CairoMakie.Axis(fig[1, 1]; xlabel="Time", ylabel="Waveform") l1 = lines!(ax, tsteps, waveform; linewidth=2, alpha=0.75) s1 = scatter!( ax, tsteps, waveform; marker=:circle, alpha=0.5, strokewidth=2, markersize=12 ) l2 = lines!(ax, tsteps, waveform_nn; linewidth=2, alpha=0.75) s2 = scatter!( ax, tsteps, waveform_nn; marker=:circle, alpha=0.5, strokewidth=2, markersize=12 ) l3 = lines!(ax, tsteps, waveform_nn_trained; linewidth=2, alpha=0.75) s3 = scatter!( ax, tsteps, waveform_nn_trained; marker=:circle, alpha=0.5, strokewidth=2, markersize=12, ) axislegend( ax, [[l1, s1], [l2, s2], [l3, s3]], ["Waveform Data", "Waveform Neural Net (Untrained)", "Waveform Neural Net"]; position=:lb, ) fig end ``` ## Appendix {#Appendix} ```julia using InteractiveUtils InteractiveUtils.versioninfo() if @isdefined(MLDataDevices) if @isdefined(CUDA) && MLDataDevices.functional(CUDADevice) println() CUDA.versioninfo() end if @isdefined(AMDGPU) && MLDataDevices.functional(AMDGPUDevice) println() AMDGPU.versioninfo() end end ``` ``` Julia Version 1.12.6 Commit 15346901f00 (2026-04-09 19:20 UTC) Build Info: Official https://julialang.org release Platform Info: OS: Linux (x86_64-linux-gnu) CPU: 4 × AMD EPYC 9V74 80-Core Processor WORD_SIZE: 64 LLVM: libLLVM-18.1.7 (ORCJIT, znver4) GC: Built with stock GC Threads: 4 default, 1 interactive, 4 GC (on 4 virtual cores) Environment: JULIA_DEBUG = Literate LD_LIBRARY_PATH = JULIA_NUM_THREADS = 4 JULIA_CPU_HARD_MEMORY_LIMIT = 100% JULIA_PKG_PRECOMPILE_AUTO = 0 ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/advanced/2_DDIM.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # Denoising Diffusion Implicit Model (DDIM) {#Denoising-Diffusion-Implicit-Model-DDIM} [Lux.jl](https://github.com/LuxDL/Lux.jl) implementation of Denoising Diffusion Implicit Models ([arXiv:2010.02502](https://arxiv.org/abs/2010.02502)). The model generates images from Gaussian noises by \denoising\ iteratively. ## Package Imports {#Package-Imports} ```julia ENV["XLA_REACTANT_GPU_MEM_FRACTION"] = get(ENV, "XLA_REACTANT_GPU_MEM_FRACTION", "0.98") using ConcreteStructs, ArgParse, DataAugmentation, DataDeps, Dates, Enzyme, FileIO, ImageCore, ImageShow, JLD2, Lux, MLUtils, Optimisers, ParameterSchedulers, ProgressTables, Printf, Random, Reactant, Statistics, TensorBoardLogger, OhMyThreads const IN_VSCODE = isdefined(Main, :VSCodeServer) ``` ## Model Definition {#Model-Definition} This DDIM implementation follows [the Keras example](https://keras.io/examples/generative/ddim/). Embed noise variances to embedding. ```julia function sinusoidal_embedding( x::AbstractArray{T,4}, min_freq, max_freq, embedding_dims::Int ) where {T} if size(x)[1:3] != (1, 1, 1) throw(DimensionMismatch("Input shape must be (1, 1, 1, batch)")) end lower, upper = T(log(min_freq)), T(log(max_freq)) n = embedding_dims ÷ 2 freqs = exp.(reshape(range(lower, upper; length=n), 1, 1, n, 1)) x_ = 2 .* x .* freqs return cat(sinpi.(x_), cospi.(x_); dims=Val(3)) end @concrete struct ResidualBlock <: AbstractLuxWrapperLayer{:layer} layer end function ResidualBlock(in_channels::Int, out_channels::Int) return ResidualBlock( Parallel( +, if in_channels == out_channels NoOpLayer() else Conv((1, 1), in_channels => out_channels; pad=SamePad()) end, Chain( BatchNorm(in_channels; affine=false), Conv((3, 3), in_channels => out_channels, swish; pad=SamePad()), Conv((3, 3), out_channels => out_channels; pad=SamePad()), ), ), ) end @concrete struct DownsampleBlock <: AbstractLuxContainerLayer{(:residual_blocks, :pool)} residual_blocks pool end function DownsampleBlock(in_channels::Int, out_channels::Int, block_depth::Int) residual_blocks = Tuple([ ResidualBlock(ifelse(i == 1, in_channels, out_channels), out_channels) for i in 1:block_depth ]) return DownsampleBlock(residual_blocks, MeanPool((2, 2))) end function (d::DownsampleBlock)(x::AbstractArray, ps, st::NamedTuple) skips = (x,) st_residual_blocks = () for i in eachindex(d.residual_blocks) y, st_new = d.residual_blocks[i]( last(skips), ps.residual_blocks[i], st.residual_blocks[i] ) skips = (skips..., y) st_residual_blocks = (st_residual_blocks..., st_new) end y, st_pool = d.pool(last(skips), ps.pool, st.pool) return (y, skips), (; residual_blocks=st_residual_blocks, pool=st_pool) end @concrete struct UpsampleBlock <: AbstractLuxContainerLayer{(:residual_blocks, :upsample)} residual_blocks upsample end function UpsampleBlock(in_channels::Int, out_channels::Int, block_depth::Int) residual_blocks = Tuple([ ResidualBlock( ifelse(i == 1, in_channels + out_channels, out_channels * 2), out_channels ) for i in 1:block_depth ]) return UpsampleBlock(residual_blocks, Upsample(:bilinear; scale=2)) end function (u::UpsampleBlock)((x, skips), ps, st::NamedTuple) x, st_upsample = u.upsample(x, ps.upsample, st.upsample) y, st_residual_blocks = x, () for i in eachindex(u.residual_blocks) y, st_new = u.residual_blocks[i]( cat(y, skips[end - i + 1]; dims=Val(3)), ps.residual_blocks[i], st.residual_blocks[i], ) st_residual_blocks = (st_residual_blocks..., st_new) end return y, (; residual_blocks=st_residual_blocks, upsample=st_upsample) end @concrete struct UNet <: AbstractLuxContainerLayer{( :conv_in, :conv_out, :down_blocks, :residual_blocks, :up_blocks, :upsample )} upsample conv_in conv_out down_blocks residual_blocks up_blocks min_freq max_freq embedding_dims end function UNet( image_size::Dims{2}; channels=[32, 64, 96, 128], block_depth=2, min_freq=1.0f0, max_freq=1000.0f0, embedding_dims=32, ) upsample = Upsample(:nearest; size=image_size) conv_in = Conv((1, 1), 3 => channels[1]) conv_out = Conv((1, 1), channels[1] => 3; init_weight=zeros32) channel_input = embedding_dims + channels[1] down_blocks = Tuple([ DownsampleBlock(i == 1 ? channel_input : channels[i - 1], channels[i], block_depth) for i in 1:(length(channels) - 1) ]) residual_blocks = Chain( [ ResidualBlock(ifelse(i == 1, channels[end - 1], channels[end]), channels[end]) for i in 1:block_depth ]..., ) reverse!(channels) up_blocks = Tuple([ UpsampleBlock(in_chs, out_chs, block_depth) for (in_chs, out_chs) in zip(channels[1:(end - 1)], channels[2:end]) ]) return UNet( upsample, conv_in, conv_out, down_blocks, residual_blocks, up_blocks, min_freq, max_freq, embedding_dims, ) end function (u::UNet)((noisy_images, noise_variances), ps, st::NamedTuple) @assert size(noise_variances)[1:3] == (1, 1, 1) @assert size(noisy_images, 4) == size(noise_variances, 4) emb, st_upsample = u.upsample( sinusoidal_embedding(noise_variances, u.min_freq, u.max_freq, u.embedding_dims), ps.upsample, st.upsample, ) tmp, st_conv_in = u.conv_in(noisy_images, ps.conv_in, st.conv_in) x = cat(tmp, emb; dims=Val(3)) skips_at_each_stage = () st_down_blocks = () for i in eachindex(u.down_blocks) (x, skips), st_new = u.down_blocks[i](x, ps.down_blocks[i], st.down_blocks[i]) skips_at_each_stage = (skips_at_each_stage..., skips) st_down_blocks = (st_down_blocks..., st_new) end x, st_residual_blocks = u.residual_blocks(x, ps.residual_blocks, st.residual_blocks) st_up_blocks = () for i in eachindex(u.up_blocks) x, st_new = u.up_blocks[i]( (x, skips_at_each_stage[end - i + 1]), ps.up_blocks[i], st.up_blocks[i] ) st_up_blocks = (st_up_blocks..., st_new) end x, st_conv_out = u.conv_out(x, ps.conv_out, st.conv_out) return ( x, (; conv_in=st_conv_in, conv_out=st_conv_out, down_blocks=st_down_blocks, residual_blocks=st_residual_blocks, up_blocks=st_up_blocks, upsample=st_upsample, ), ) end function diffusion_schedules( diffusion_times::AbstractArray{T,4}, min_signal_rate, max_signal_rate ) where {T} start_angle = T(acos(max_signal_rate)) end_angle = T(acos(min_signal_rate)) diffusion_angles = @. start_angle + (end_angle - start_angle) * diffusion_times signal_rates = @. cos(diffusion_angles) noise_rates = @. sin(diffusion_angles) return noise_rates, signal_rates end function denoise( unet, noisy_images::AbstractArray{T1,4}, noise_rates::AbstractArray{T2,4}, signal_rates::AbstractArray{T3,4}, ) where {T1,T2,T3} T = promote_type(T1, T2, T3) noisy_images = T.(noisy_images) signal_rates = T.(signal_rates) pred_noises = unet((noisy_images, noise_rates .^ 2)) pred_images = @. (noisy_images - pred_noises * noise_rates) / signal_rates return pred_noises, pred_images end function denoise!( pred_images, unet, noisy_images::AbstractArray{T1,4}, noise_rates::AbstractArray{T2,4}, signal_rates::AbstractArray{T3,4}, ) where {T1,T2,T3} T = promote_type(T1, T2, T3) noisy_images = T.(noisy_images) noise_rates = T.(noise_rates) signal_rates = T.(signal_rates) pred_noises = unet((noisy_images, noise_rates .^ 2)) @. pred_images = (noisy_images - pred_noises * noise_rates) / signal_rates return pred_noises end @concrete struct DDIM <: AbstractLuxContainerLayer{(:unet, :bn)} unet bn min_signal_rate max_signal_rate image_size::Dims{3} end function DDIM( image_size::Dims{2}, args...; min_signal_rate=0.02f0, max_signal_rate=0.95f0, kwargs... ) return DDIM( UNet(image_size, args...; kwargs...), BatchNorm(3; affine=false, track_stats=true), min_signal_rate, max_signal_rate, (image_size..., 3), ) end function Lux.initialstates(rng::AbstractRNG, ddim::DDIM) rand(rng, 1) return (; rng, bn=Lux.initialstates(rng, ddim.bn), unet=Lux.initialstates(rng, ddim.unet) ) end function (ddim::DDIM)(x::AbstractArray{T,4}, ps, st::NamedTuple) where {T} images, st_bn = ddim.bn(x, ps.bn, st.bn) rng = Lux.replicate(st.rng) diffusion_times = rand_like(rng, images, (1, 1, 1, size(images, 4))) noise_rates, signal_rates = diffusion_schedules( diffusion_times, ddim.min_signal_rate, ddim.max_signal_rate ) noises = randn_like(rng, images) noisy_images = @. signal_rates * images + noise_rates * noises unet = StatefulLuxLayer{true}(ddim.unet, ps.unet, st.unet) pred_noises, pred_images = denoise(unet, noisy_images, noise_rates, signal_rates) return ((noises, images, pred_noises, pred_images), (; rng, bn=st_bn, unet=unet.st)) end # Helper Functions for Image Generation function generate(model::DDIM, ps, st::NamedTuple, diffusion_steps::Int, num_samples::Int) rng = Lux.replicate(st.rng) μ, σ² = st.bn.running_mean, st.bn.running_var initial_noise = randn_like(rng, μ, (model.image_size..., num_samples)) generated_images = reverse_diffusion(model, initial_noise, ps, st, diffusion_steps) return clamp01.(denormalize(generated_images, μ, σ², model.bn.epsilon)) end function reverse_diffusion_single_step!( pred_images, noisy_images, step, step_size, unet, ps, st, ones, min_signal_rate, max_signal_rate, ) diffusion_times = ones .- step_size * step noise_rates, signal_rates = diffusion_schedules( diffusion_times, min_signal_rate, max_signal_rate ) sunet = StatefulLuxLayer{true}(unet, ps, st) pred_noises = denoise!(pred_images, sunet, noisy_images, noise_rates, signal_rates) next_diffusion_times = diffusion_times .- step_size next_noisy_rates, next_signal_rates = diffusion_schedules( next_diffusion_times, min_signal_rate, max_signal_rate ) @. noisy_images = next_signal_rates .* pred_images .+ next_noisy_rates .* pred_noises return nothing end function reverse_diffusion( model::DDIM, initial_noise::AbstractArray{T,4}, ps, st::NamedTuple, diffusion_steps::Int ) where {T} step_size = one(T) / diffusion_steps ones_dev = ones_like(initial_noise, (1, 1, 1, size(initial_noise, 4))) noisy_images, pred_images = initial_noise, similar(initial_noise) @trace for step in 1:diffusion_steps reverse_diffusion_single_step!( pred_images, noisy_images, step, step_size, model.unet, ps.unet, st.unet, ones_dev, model.min_signal_rate, model.max_signal_rate, ) end return pred_images end function denormalize(x::AbstractArray{T,4}, μ, σ², ε) where {T} μ = reshape(μ, 1, 1, 3, 1) σ = sqrt.(reshape(σ², 1, 1, 3, 1) .+ T(ε)) return σ .* x .+ μ end function create_image_list(imgs::AbstractArray) return map(eachslice(imgs; dims=4)) do img cimg = if size(img, 3) == 1 colorview(Gray, view(img, :, :, 1)) else colorview(RGB, permutedims(img, (3, 1, 2))) end return cimg' end end function create_image_grid( images::AbstractArray, grid_rows::Int, grid_cols::Union{Int,Nothing}=nothing ) images = ndims(images) != 1 ? create_image_list(images) : images grid_cols = grid_cols === nothing ? length(images) ÷ grid_rows : grid_cols # Check if the number of images matches the grid total_images = grid_rows * grid_cols @assert length(images) ≤ total_images # Get the size of a single image (assuming all images are the same size) img_height, img_width = size(images[1]) # Create a blank grid canvas grid_height = img_height * grid_rows grid_width = img_width * grid_cols grid_canvas = similar(images[1], grid_height, grid_width) # Place each image in the correct position on the canvas for idx in 1:total_images idx > length(images) && break row = div(idx - 1, grid_cols) + 1 col = mod(idx - 1, grid_cols) + 1 start_row = (row - 1) * img_height + 1 start_col = (col - 1) * img_width + 1 grid_canvas[start_row:(start_row + img_height - 1), start_col:(start_col + img_width - 1)] .= images[idx] end return grid_canvas end function save_images(output_dir, images::Vector{<:AbstractMatrix{<:RGB}}) for (i, img) in enumerate(images) if any(isnan, img) @warn "NaN image found in the generated images. Skipping..." continue end save(joinpath(output_dir, "img_$(i).png"), img) end end ``` ## Dataset {#Dataset} We will register the dataset using the [DataDeps.jl](https://github.com/oxinabox/DataDeps.jl) package. The dataset is available at . This allows us to automatically download the dataset when we run the code. ```julia @concrete struct FlowersDataset image_size image_files transform end FlowersDataset(image_size::Int) = FlowersDataset((image_size, image_size)) function FlowersDataset(image_size::Dims{2}) dirpath = try joinpath(datadep"FlowersDataset", "jpg") catch err err isa KeyError || rethrow() register( DataDep( "FlowersDataset", "102 Category Flowers Dataset", "https://www.robots.ox.ac.uk/~vgg/data/flowers/102/102flowers.tgz", "2d01ecc807db462958cfe3d92f57a8c252b4abd240eb955770201e45f783b246"; post_fetch_method=file -> run(`tar -xzf $file`), ), ) joinpath(datadep"FlowersDataset", "jpg") end image_files = joinpath.(dirpath, readdir(dirpath)) transform = ScaleKeepAspect(image_size) |> CenterResizeCrop(image_size) |> Maybe(FlipX{2}()) |> ImageToTensor() return FlowersDataset(image_size, image_files, transform) end Base.length(ds::FlowersDataset) = length(ds.image_files) function Base.getindex(ds::FlowersDataset, i::Int) return Float32.(itemdata(apply(ds.transform, Image(load(ds.image_files[i]))))) end function Base.getindex(ds::FlowersDataset, idxs) imgs = Array{Float32,4}(undef, ds.image_size..., 3, length(idxs)) tforeach(1:length(idxs)) do i img = Image(load(ds.image_files[idxs[i]])) return copyto!(view(imgs, :, :, :, i), itemdata(apply(ds.transform, img))) end return imgs end function loss_function(model, ps, st, data) (noises, images, pred_noises, pred_images), stₙ = Lux.apply(model, data, ps, st) noise_loss = MSELoss()(pred_noises, noises) image_loss = MSELoss()(pred_images, images) return noise_loss, stₙ, (; image_loss, noise_loss) end ``` ## Entry Point for our code {#Entry-Point-for-our-code} ```julia function main(; epochs::Int=100, image_size::Int=128, batchsize::Int=128, learning_rate_start::Float32=1.0f-3, learning_rate_end::Float32=1.0f-5, weight_decay::Float32=1.0f-6, checkpoint_interval::Int=25, expt_dir="", diffusion_steps::Int=80, generate_image_interval::Int=5, # model hyper params channels::Vector{Int}=[32, 64, 96, 128], block_depth::Int=2, min_freq::Float32=1.0f0, max_freq::Float32=1000.0f0, embedding_dims::Int=32, min_signal_rate::Float32=0.02f0, max_signal_rate::Float32=0.95f0, # inference specific inference_mode::Bool=false, saved_model_path=nothing, generate_n_images::Int=12, ) if isempty(expt_dir) expt_dir = joinpath(@__DIR__, string(now(UTC)) * "_" * uppercase(randstring(4))) * "_ddim" end isdir(expt_dir) || mkpath(expt_dir) @printf "[%s] [Info] Experiment directory: %s\n" now(UTC) expt_dir rng = Random.default_rng() Random.seed!(rng, 1234) image_dir = joinpath(expt_dir, "images") isdir(image_dir) || mkpath(image_dir) ckpt_dir = joinpath(expt_dir, "checkpoints") isdir(ckpt_dir) || mkpath(ckpt_dir) xdev = reactant_device(; force=true) cdev = cpu_device() @printf "[%s] [Info] Building model\n" now(UTC) model = DDIM( (image_size, image_size); channels, block_depth, min_freq, max_freq, embedding_dims, min_signal_rate, max_signal_rate, ) if inference_mode @assert saved_model_path !== nothing "`saved_model_path` must be specified for \ inference" @load saved_model_path parameters states ps, st = (parameters, states) |> xdev generated_images = @jit generate( model, ps, Lux.testmode(st), diffusion_steps, generate_n_images ) generated_images = generated_images |> cdev path = joinpath(image_dir, "inference") @printf "[%s] [Info] Saving generated images to %s\n" now(UTC) path imgs = create_image_list(generated_images) save_images(path, imgs) if IN_VSCODE display(create_image_grid(generated_images, 8, cld(length(imgs), 8))) end return nothing end ps, st = Lux.setup(rng, model) |> xdev tb_dir = joinpath(expt_dir, "tb_logs") @printf "[%s] [Info] Tensorboard logs being saved to %s. Run tensorboard with \ `tensorboard --logdir %s`\n" now(UTC) tb_dir dirname(tb_dir) tb_logger = TBLogger(tb_dir) opt = AdamW(; eta=learning_rate_start, lambda=weight_decay) scheduler = CosAnneal(learning_rate_start, learning_rate_end, epochs) tstate = Training.TrainState(model, ps, st, opt) @printf "[%s] [Info] Preparing dataset\n" now(UTC) ds = FlowersDataset(image_size) data_loader = DataLoader(ds; batchsize, shuffle=true, partial=false, parallel=false) |> xdev pt = ProgressTable(; header=["Epoch", "Image Loss", "Noise Loss", "Time (s)", "Throughput (img/s)"], widths=[10, 24, 24, 24, 24], format=["%3d", "%.6f", "%.6f", "%.6f", "%.6f"], color=[:normal, :normal, :normal, :normal, :normal], border=true, alignment=[:center, :center, :center, :center, :center], ) @printf "[%s] [Info] Compiling generate function\n" now(UTC) time_start = time() generate_compiled = @compile generate( model, ps, Lux.testmode(st), diffusion_steps, generate_n_images ) @printf "[%s] [Info] Compiled generate function in %.6f seconds\n" now(UTC) ( time() - time_start ) image_losses = Vector{Float32}(undef, length(data_loader)) noise_losses = Vector{Float32}(undef, length(data_loader)) step = 1 @printf "[%s] [Info] Training model\n" now(UTC) initialize(pt) for epoch in 1:epochs total_time = 0.0 total_samples = 0 eta = Float32(scheduler(epoch)) tstate = Optimisers.adjust!(tstate, eta) log_value(tb_logger, "Training/Learning Rate", eta; step) start_time = time() for (i, data) in enumerate(data_loader) (_, loss, stats, tstate) = Training.single_train_step!( AutoEnzyme(), loss_function, data, tstate; return_gradients=Val(false) ) @assert !isnan(loss) "NaN loss ($(loss)) encountered!" total_samples += size(data, ndims(data)) image_losses[i] = stats.image_loss noise_losses[i] = stats.noise_loss log_value(tb_logger, "Training/Image Loss", Float32(stats.image_loss); step) log_value(tb_logger, "Training/Noise Loss", Float32(stats.noise_loss); step) log_value( tb_logger, "Training/Throughput", total_samples / (time() - start_time); step, ) step += 1 end total_time = time() - start_time next( pt, [ epoch, mean(image_losses), mean(noise_losses), total_time, total_samples / total_time, ], ) if epoch % generate_image_interval == 0 || epoch == epochs generated_images = generate_compiled( tstate.model, tstate.parameters, Lux.testmode(tstate.states), diffusion_steps, generate_n_images, ) generated_images = generated_images |> cdev path = joinpath(image_dir, "epoch_$(epoch)") imgs = create_image_list(generated_images) save_images(path, imgs) log_images(tb_logger, "Generated Images", imgs; step) if IN_VSCODE display(create_image_grid(generated_images, 8, cld(length(imgs), 8))) end end if epoch % checkpoint_interval == 0 || epoch == epochs path = joinpath(ckpt_dir, "model_$(epoch).jld2") @printf "[%s] [Info] Saving checkpoint to %s\n" now(UTC) path parameters = tstate.parameters |> cdev states = tstate.states |> cdev @save path parameters states end end finalize(pt) @printf "[%s] [Info] Saving final model\n" now(UTC) return tstate end function get_argparse_settings() s = ArgParseSettings(; autofix_names=true) #! format: off @add_arg_table s begin "--epochs" help = "Number of epochs to train" arg_type = Int default = 100 "--image-size" help = "Input image size (square)" arg_type = Int default = 128 "--batchsize" help = "Training batch size" arg_type = Int default = 128 "--learning-rate-start" help = "Starting learning rate" arg_type = Float32 default = 3.0f-3 "--learning-rate-end" help = "Final learning rate" arg_type = Float32 default = 1.0f-4 "--weight-decay" help = "Weight decay (AdamW lambda)" arg_type = Float32 default = 1.0f-6 "--checkpoint-interval" help = "Save checkpoint every N epochs" arg_type = Int default = 25 "--expt-dir" help = "Experiment output directory" arg_type = String default = "" "--diffusion-steps" help = "Number of DDIM reverse diffusion steps" arg_type = Int default = 80 "--generate-image-interval" help = "Generate and log images every N epochs" arg_type = Int default = 5 ``` model hyper params ```julia "--channels" help = "UNet channels per stage" arg_type = Int nargs = '+' default = [32, 64, 96, 128] "--block-depth" help = "Number of residual blocks per stage" arg_type = Int default = 2 "--min-freq" help = "Sinusoidal embedding min frequency" arg_type = Float32 default = 1.0f0 "--max-freq" help = "Sinusoidal embedding max frequency" arg_type = Float32 default = 1000.0f0 "--embedding-dims" help = "Sinusoidal embedding dimension" arg_type = Int default = 32 "--min-signal-rate" help = "Minimum signal rate" arg_type = Float32 default = 0.02f0 "--max-signal-rate" help = "Maximum signal rate" arg_type = Float32 default = 0.95f0 ``` inference specific ```julia "--inference" help = "Run in inference-only mode" action = :store_true "--saved-model-path" help = "Path to JLD2 checkpoint (required with --inference)" arg_type = String "--generate-n-images" help = "Number of images to generate during inference or periodic logging" arg_type = Int default = 12 end #! format: on return s end if abspath(PROGRAM_FILE) == @__FILE__ args = parse_args(ARGS, get_argparse_settings(); as_symbols=true) main(; epochs=args[:epochs], image_size=args[:image_size], batchsize=args[:batchsize], learning_rate_start=args[:learning_rate_start], learning_rate_end=args[:learning_rate_end], weight_decay=args[:weight_decay], checkpoint_interval=args[:checkpoint_interval], expt_dir=args[:expt_dir], diffusion_steps=args[:diffusion_steps], generate_image_interval=args[:generate_image_interval], channels=args[:channels], block_depth=args[:block_depth], min_freq=args[:min_freq], max_freq=args[:max_freq], embedding_dims=args[:embedding_dims], min_signal_rate=args[:min_signal_rate], max_signal_rate=args[:max_signal_rate], inference_mode=args[:inference], saved_model_path=args[:saved_model_path], generate_n_images=args[:generate_n_images], ) end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/advanced/3_ImageNet.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # ImageNet Classification using Distributed Data Parallel Training {#ImageNet-Classification-using-Distributed-Data-Parallel-Training} This implements training of popular model architectures, such as ResNet, AlexNet, and VGG on the ImageNet dataset. For distributed data-parallel training we need to launch this script using `mpiexecjl` Setup [MPI.jl](https://juliaparallel.org/MPI.jl/). If your system has functional NCCL we will use it for all CUDA communications. Otherwise, we will use MPI for all communications. ```bash mpiexecjl -np 4 julia --startup=no --project=examples/ImageNet -t auto\ examples/ImageNet/main.jl \ --model-name="ViT" \ --model-kind="tiny" \ --train-batchsize=256 \ --val-batchsize=256 \ --optimizer-kind="sgd" \ --learning-rate=0.01 \ --base-path="/home/avik-pal/data/ImageNet/" ``` For single-node training, we can simply launch the script using `julia` ```bash julia --startup=no --project=examples/ImageNet -t auto examples/ImageNet/main.jl \ --model-name="ViT" \ --model-kind="tiny" \ --train-batchsize=256 \ --val-batchsize=256 \ --optimizer-kind="sgd" \ --learning-rate=0.01 \ --base-path="/home/avik-pal/data/ImageNet/" ``` ## Package Imports {#Package-Imports} ```julia using Boltz, Lux, MLDataDevices # import Metalhead # Install and load this package to use the Metalhead models with Lux using Dates, Random using DataAugmentation, FileIO, MLUtils, OneHotArrays, Optimisers, ParameterSchedulers, Setfield using Comonicon, Format using JLD2 using Zygote using LuxCUDA # using AMDGPU # Install and load AMDGPU to train models on AMD GPUs with ROCm using MPI: MPI # Enables distributed training in Lux. NCCL is needed for CUDA GPUs using NCCL: NCCL const gdev = gpu_device() const cdev = cpu_device() ``` ## Setup Distributed Training {#Setup-Distributed-Training} We will use NCCL for NVIDIA GPUs and MPI for anything else ```julia const distributed_backend = try if gdev isa CUDADevice DistributedUtils.initialize(NCCLBackend) DistributedUtils.get_distributed_backend(NCCLBackend) else DistributedUtils.initialize(MPIBackend) DistributedUtils.get_distributed_backend(MPIBackend) end catch err @error "Could not initialize distributed training. Error: $err" nothing end const local_rank = distributed_backend === nothing ? 0 : DistributedUtils.local_rank(distributed_backend) const total_workers = if distributed_backend === nothing 1 else DistributedUtils.total_workers(distributed_backend) end const is_distributed = total_workers > 1 const should_log = !is_distributed || local_rank == 0 ``` ## Data Loading for ImageNet {#Data-Loading-for-ImageNet} ```julia # We need the data to be in a specific format. See the # [README.md](https://github.com/LuxDL/Lux.jl/blob/main/examples/ImageNet/README.md) for more details. const IMAGENET_CORRUPTED_FILES = [ "n01739381_1309.JPEG", "n02077923_14822.JPEG", "n02447366_23489.JPEG", "n02492035_15739.JPEG", "n02747177_10752.JPEG", "n03018349_4028.JPEG", "n03062245_4620.JPEG", "n03347037_9675.JPEG", "n03467068_12171.JPEG", "n03529860_11437.JPEG", "n03544143_17228.JPEG", "n03633091_5218.JPEG", "n03710637_5125.JPEG", "n03961711_5286.JPEG", "n04033995_2932.JPEG", "n04258138_17003.JPEG", "n04264628_27969.JPEG", "n04336792_7448.JPEG", "n04371774_5854.JPEG", "n04596742_4225.JPEG", "n07583066_647.JPEG", "n13037406_4650.JPEG", "n02105855_2933.JPEG", "ILSVRC2012_val_00019877.JPEG", ] function load_imagenet1k(base_path::String, split::Symbol) @assert split in (:train, :val) full_path = joinpath(base_path, string(split)) synsets = sort(readdir(full_path)) @assert length(synsets) == 1000 "There should be 1000 subdirectories in $(full_path)." image_files = String[] labels = Int[] for (i, synset) in enumerate(synsets) filenames = readdir(joinpath(full_path, synset)) filter!(x -> x ∉ IMAGENET_CORRUPTED_FILES, filenames) paths = joinpath.((full_path,), (synset,), filenames) append!(image_files, paths) append!(labels, repeat([i - 1], length(paths))) end return image_files, labels end default_image_size(::Type{Vision.VisionTransformer}, ::Nothing) = 256 default_image_size(::Type{Vision.VisionTransformer}, size::Int) = size default_image_size(_, ::Nothing) = 224 default_image_size(_, size::Int) = size struct MakeColoredImage <: DataAugmentation.Transform end function DataAugmentation.apply( ::MakeColoredImage, item::DataAugmentation.AbstractArrayItem; randstate=nothing ) data = itemdata(item) (ndims(data) == 2 || size(data, 3) == 1) && (data = cat(data, data, data; dims=Val(3))) return DataAugmentation.setdata(item, data) end struct FileDataset files labels augment end Base.length(dataset::FileDataset) = length(dataset.files) function Base.getindex(dataset::FileDataset, i::Int) img = Image(FileIO.load(dataset.files[i])) aug_img = itemdata(DataAugmentation.apply(dataset.augment, img)) return aug_img, OneHotArrays.onehot(dataset.labels[i], 0:999) end function construct_dataloaders(; base_path::String, train_batchsize, val_batchsize, image_size::Int ) sensible_println("=> creating dataloaders.") train_augment = ScaleFixed((256, 256)) |> Maybe(FlipX(), 0.5) |> RandomResizeCrop((image_size, image_size)) |> PinOrigin() |> ImageToTensor() |> MakeColoredImage() |> ToEltype(Float32) |> Normalize((0.485f0, 0.456f0, 0.406f0), (0.229f0, 0.224f0, 0.225f0)) train_files, train_labels = load_imagenet1k(base_path, :train) train_dataset = FileDataset(train_files, train_labels, train_augment) val_augment = ScaleFixed((image_size, image_size)) |> PinOrigin() |> ImageToTensor() |> MakeColoredImage() |> ToEltype(Float32) |> Normalize((0.485f0, 0.456f0, 0.406f0), (0.229f0, 0.224f0, 0.225f0)) val_files, val_labels = load_imagenet1k(base_path, :val) val_dataset = FileDataset(val_files, val_labels, val_augment) if is_distributed train_dataset = DistributedUtils.DistributedDataContainer( distributed_backend, train_dataset ) val_dataset = DistributedUtils.DistributedDataContainer( distributed_backend, val_dataset ) end train_dataloader = DataLoader( train_dataset; batchsize=train_batchsize ÷ total_workers, partial=false, collate=true, shuffle=true, parallel=true, ) val_dataloader = DataLoader( val_dataset; batchsize=val_batchsize ÷ total_workers, partial=true, collate=true, shuffle=false, parallel=true, ) return gdev(train_dataloader), gdev(val_dataloader) end ``` ## Model Construction {#Model-Construction} ```julia function construct_model(; rng::AbstractRNG, model_name::String, model_args, pretrained::Bool=false ) model = getproperty(Vision, Symbol(model_name))(model_args...; pretrained) ps, st = Lux.setup(rng, model) |> gdev sensible_println("=> model `$(model_name)` created.") pretrained && sensible_println("==> using pre-trained model`") sensible_println("==> number of trainable parameters: $(Lux.parameterlength(ps))") sensible_println("==> number of states: $(Lux.statelength(st))") if is_distributed ps = DistributedUtils.synchronize!!(distributed_backend, ps) st = DistributedUtils.synchronize!!(distributed_backend, st) sensible_println("==> synced model parameters and states across all ranks") end return model, ps, st end ``` ## Optimizer Configuration {#Optimizer-Configuration} ```julia function construct_optimizer_and_scheduler(; kind::String, learning_rate::AbstractFloat, nesterov::Bool, momentum::AbstractFloat, weight_decay::AbstractFloat, scheduler_kind::String, cycle_length::Int, damp_factor::AbstractFloat, lr_step_decay::AbstractFloat, lr_step::Vector{Int}, ) sensible_println("=> creating optimizer.") kind = Symbol(kind) optimizer = if kind == :adam Adam(learning_rate) elseif kind == :sgd if nesterov Nesterov(learning_rate, momentum) elseif iszero(momentum) Descent(learning_rate) else Momentum(learning_rate, momentum) end else throw(ArgumentError("Unknown value for `optimizer` = $kind. Supported options are: \ `adam` and `sgd`.")) end optimizer = if iszero(weight_decay) optimizer else OptimiserChain(optimizer, WeightDecay(weight_decay)) end sensible_println("=> creating scheduler.") scheduler_kind = Symbol(scheduler_kind) scheduler = if scheduler_kind == :cosine l0 = learning_rate l1 = learning_rate / 100 ComposedSchedule( CosAnneal(l0, l1, cycle_length), Step(l0, damp_factor, cycle_length) ) elseif scheduler_kind == :constant Constant(learning_rate) elseif scheduler_kind == :step Step(learning_rate, lr_step_decay, lr_step) else throw(ArgumentError("Unknown value for `lr_scheduler` = $(scheduler_kind). \ Supported options are: `constant`, `step` and `cosine`.")) end optimizer = if is_distributed DistributedUtils.DistributedOptimizer(distributed_backend, optimizer) else optimizer end return optimizer, scheduler end ``` ## Utility Functions {#Utility-Functions} ```julia const logitcrossentropy = CrossEntropyLoss(; logits=Val(true)) function loss_function(model, ps, st, (img, y)) ŷ, stₙ = model(img, ps, st) return logitcrossentropy(ŷ, y), stₙ, (; prediction=ŷ) end sensible_println(msg) = should_log && println("[$(now())] ", msg) sensible_print(msg) = should_log && print("[$(now())] ", msg) function accuracy(ŷ::AbstractMatrix, y::AbstractMatrix, topk=(1,)) pred_labels = partialsortperm.(eachcol(cdev(ŷ)), Ref(1:maximum(topk)); rev=true) true_labels = onecold(cdev(y)) accuracies = Vector{Float64}(undef, length(topk)) for (i, k) in enumerate(topk) accuracies[i] = sum( map((a, b) -> sum(view(a, 1:k) .== b), pred_labels, true_labels) ) end accuracies .= accuracies .* 100 ./ size(y, 2) return accuracies end function save_checkpoint(state::NamedTuple; is_best::Bool, filename::String) should_log || return nothing @assert last(splitext(filename)) == ".jld2" "Filename should have a .jld2 extension." isdir(dirname(filename)) || mkpath(dirname(filename)) save(filename; state) sensible_println("=> saved checkpoint `$(filename)`.") if is_best symlink_safe(filename, joinpath(dirname(filename), "model_best.jld2")) sensible_println("=> best model updated to `$(filename)`!") end symlink_safe(filename, joinpath(dirname(filename), "model_current.jld2")) return nothing end function symlink_safe(src, dest) rm(dest; force=true) symlink(src, dest) return nothing end function load_checkpoint(filename::String) try ## NOTE(@avik-pal): ispath is failing for symlinks? return JLD2.load(filename)[:state] catch sensible_println("$(filename) could not be loaded. This might be because the file \ is absent or is corrupt. Proceeding by returning `nothing`.") return nothing end end function full_gc_and_reclaim() GC.gc(true) MLDataDevices.functional(CUDADevice) && CUDA.reclaim() MLDataDevices.functional(AMDGPUDevice) && AMDGPU.reclaim() return nothing end @kwdef mutable struct AverageMeter fmtstr val::Float64 = 0.0 sum::Float64 = 0.0 count::Int = 0 average::Float64 = 0 end function AverageMeter(name::String, fmt::String) return AverageMeter(; fmtstr=FormatExpr("$(name) {1:$(fmt)} ({2:$(fmt)})")) end function (meter::AverageMeter)(val, n::Int) meter.val = val s = val * n if is_distributed v = [s, typeof(val)(n)] DistributedUtils.allreduce!(backend, v, +) s, n = v[1], Int(v[2]) end meter.sum += s meter.count += n meter.average = meter.sum / meter.count return meter.average end function reset_meter!(meter::AverageMeter) meter.val = 0.0 meter.sum = 0.0 meter.count = 0 meter.average = 0.0 return meter end function print_meter(meter::AverageMeter) return should_log && printfmt(meter.fmtstr, meter.val, meter.average) end struct ProgressMeter batch_fmtstr meters end function ProgressMeter(num_batches::Int, meters, prefix::String="") fmt = "%" * string(length(string(num_batches))) * "d" fmt2 = "{1:" * string(length(string(num_batches))) * "d}" prefix = prefix != "" ? endswith(prefix, " ") ? prefix : prefix * " " : "" batch_fmtstr = FormatExpr("$prefix[$fmt2/" * cfmt(fmt, num_batches) * "]") return ProgressMeter(batch_fmtstr, meters) end reset_meter!(meter::ProgressMeter) = foreach(reset_meter!, meter.meters) function print_meter(meter::ProgressMeter, batch::Int) should_log || return nothing printfmt(meter.batch_fmtstr, batch) foreach(meter.meters) do x print("\t") print_meter(x) return nothing end println() return nothing end get_loggable_values(meter::ProgressMeter) = getproperty.(meter.meters, :average) ``` ## Training and Validation Loops {#Training-and-Validation-Loops} ```julia function validate(val_loader, model, ps, st, step, total_steps) batch_time = AverageMeter("Batch Time", "6.5f") data_time = AverageMeter("Data Time", "6.5f") forward_time = AverageMeter("Forward Pass Time", "6.5f") losses = AverageMeter("Loss", ".6f") top1 = AverageMeter("Acc@1", "6.4f") top5 = AverageMeter("Acc@5", "6.4f") progress = ProgressMeter( total_steps, (batch_time, data_time, forward_time, losses, top1, top5), "Val:" ) st = Lux.testmode(st) t = time() for (img, y) in val_loader t_data, t = time() - t, time() bsize = size(img, ndims(img)) loss, st, stats = loss_function(model, ps, st, (img, y)) t_forward = time() - t acc1, acc5 = accuracy(stats.prediction, y, (1, 5)) top1(acc1, bsize) top5(acc5, bsize) losses(loss, bsize) data_time(t_data, bsize) forward_time(t_forward, bsize) batch_time(t_data + t_forward, bsize) t = time() end print_meter(progress, step) return top1.average end ``` ## Entry Point {#Entry-Point} ```julia Comonicon.@main function main(; seed::Int=0, model_name::String, model_kind::String="nokind", depth::Int=-1, pretrained::Bool=false, base_path::String="", train_batchsize::Int=64, val_batchsize::Int=64, image_size::Int=-1, optimizer_kind::String="sgd", learning_rate::Float32=0.01f0, nesterov::Bool=false, momentum::Float32=0.0f0, weight_decay::Float32=0.0f0, scheduler_kind::String="step", cycle_length::Int=50000, damp_factor::Float32=1.2f0, lr_step_decay::Float32=0.1f0, lr_step::Vector{Int}=[100000, 250000, 500000], expt_id::String="", expt_subdir::String=@__DIR__, resume::String="", evaluate::Bool=false, total_steps::Int=800000, evaluate_every::Int=10000, print_frequency::Int=100, ) best_acc1 = 0 rng = Random.default_rng() Random.seed!(rng, seed) model_type = getproperty(Vision, Symbol(model_name)) image_size = default_image_size(model_type, image_size == -1 ? nothing : image_size) depth = depth == -1 ? nothing : depth model_kind = model_kind == "nokind" ? nothing : Symbol(model_kind) model_args = if model_kind === nothing && depth === nothing () elseif model_kind !== nothing (model_kind,) else (depth,) end model, ps, st = construct_model(; rng, model_name, model_args, pretrained) ds_train, ds_val = construct_dataloaders(; base_path, train_batchsize, val_batchsize, image_size ) opt, scheduler = construct_optimizer_and_scheduler(; kind=optimizer_kind, learning_rate, nesterov, momentum, weight_decay, scheduler_kind, cycle_length, damp_factor, lr_step_decay, lr_step, ) expt_name = "name-$(model_name)_seed-$(seed)_id-$(expt_id)" ckpt_dir = joinpath(expt_subdir, "checkpoints", expt_name) rpath = resume == "" ? joinpath(ckpt_dir, "model_current.jld2") : resume ckpt = load_checkpoint(rpath) if !isnothing(ckpt) ps, st = (ckpt.ps, ckpt.st) |> gdev initial_step = ckpt.step sensible_println("=> training started from $(initial_step)") else initial_step = 1 end validate(ds_val, model, ps, st, 0, total_steps) evaluate && return nothing full_gc_and_reclaim() batch_time = AverageMeter("Batch Time", "6.5f") data_time = AverageMeter("Data Time", "6.5f") training_time = AverageMeter("Training Time", "6.5f") losses = AverageMeter("Loss", ".6f") top1 = AverageMeter("Acc@1", "6.4f") top5 = AverageMeter("Acc@5", "6.4f") progress = ProgressMeter( total_steps, (batch_time, data_time, training_time, losses, top1, top5), "Train:" ) st = Lux.trainmode(st) train_state = Training.TrainState(model, ps, st, opt) if is_distributed @set! train_state.optimizer_state = DistributedUtils.synchronize!!( distributed_backend, train_state.optimizer_state ) end train_loader = Iterators.cycle(ds_train) _, train_loader_state = iterate(train_loader) for step in initial_step:total_steps t = time() (img, y), train_loader_state = iterate(train_loader, train_loader_state) t_data = time() - t bsize = size(img, ndims(img)) t = time() _, loss, stats, train_state = Training.single_train_step!( AutoZygote(), loss_function, (img, y), train_state ) t_training = time() - t isnan(loss) && throw(ArgumentError("NaN loss encountered.")) acc1, acc5 = accuracy(stats.prediction, y, (1, 5)) top1(acc1, bsize) top5(acc5, bsize) losses(loss, bsize) data_time(t_data, bsize) training_time(t_training, bsize) batch_time(t_data + t_training, bsize) if step % print_frequency == 1 || step == total_steps print_meter(progress, step) reset_meter!(progress) end if step % evaluate_every == 0 acc1 = validate(ds_val, model, ps, st, step, total_steps) is_best = acc1 > best_acc1 best_acc1 = max(acc1, best_acc1) save_state = (; ps=cdev(ps), st=cdev(st), step) if should_log() save_checkpoint( save_state; is_best, filename=joinpath(ckpt_dir, "model_$(step).jld2") ) end end Optimisers.adjust!(train_state.optimizer_state, scheduler(step + 1)) end return nothing end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/tutorials/advanced/4_Qwen3.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # Qwen3 Implementation from Scratch {#Qwen3-Implementation-from-Scratch} This is an implementation of Qwen 3 ([blog](https://qwenlm.github.io/blog/qwen3/) and [technical report](https://arxiv.org/abs/2505.09388)) from scratch based on the pytorch [implementation](https://github.com/rasbt/LLMs-from-scratch/blob/main/ch05/11_qwen3/standalone-qwen3.ipynb) [developed in Pytorch under the Apache License 2.0](https://github.com/rasbt/LLMs-from-scratch/blob/main/LICENSE.txt). ## Package Imports {#Package-Imports} ```julia using BFloat16s, ConcreteStructs, LinearAlgebra, Lux, Random, Reactant using HuggingFaceTokenizers, PythonCall, SafeTensors, Scratch, JSON3 const huggingface_hub = pyimport("huggingface_hub") ``` ## Qwen3 Configuration {#Qwen3-Configuration} ```julia @kwdef struct Qwen3Config{F} version::String vocab_size::Int context_length::Int emb_dim::Int n_heads::Int n_layers::Int hidden_dim::Int head_dim::Int n_kv_groups::Int rope_base::Float32 dtype::F reasoning_model::Bool = true end function Qwen3Config(version::String; kwargs...) if version == "0.6B" return Qwen3Config(; version, vocab_size=151_936, context_length=40_960, emb_dim=1024, n_heads=16, n_layers=28, hidden_dim=3072, head_dim=128, n_kv_groups=8, rope_base=1.0f6, dtype=bf16, kwargs..., ) elseif version == "1.7B" return Qwen3Config(; version, vocab_size=151_936, context_length=40_960, emb_dim=2048, n_heads=16, n_layers=28, hidden_dim=6144, head_dim=128, n_kv_groups=8, rope_base=1.0f6, dtype=bf16, kwargs..., ) elseif version == "4B" return Qwen3Config(; version, vocab_size=151_936, context_length=40_960, emb_dim=2560, n_heads=32, n_layers=36, hidden_dim=9728, head_dim=128, n_kv_groups=8, rope_base=1.0f6, dtype=bf16, kwargs..., ) elseif version == "8B" return Qwen3Config(; version, vocab_size=151_936, context_length=40_960, emb_dim=4096, n_heads=32, n_layers=36, hidden_dim=12288, head_dim=128, n_kv_groups=8, rope_base=1.0f6, dtype=bf16, kwargs..., ) elseif version == "14B" return Qwen3Config(; version, vocab_size=151_936, context_length=40_960, emb_dim=5120, n_heads=40, n_layers=40, hidden_dim=17408, head_dim=128, n_kv_groups=8, rope_base=1.0f6, dtype=bf16, kwargs..., ) elseif version == "32B" return Qwen3Config(; version, vocab_size=151_936, context_length=40_960, emb_dim=5120, n_heads=64, n_layers=64, hidden_dim=25600, head_dim=128, n_kv_groups=8, rope_base=1.0f6, dtype=bf16, kwargs..., ) end throw(ArgumentError("Unknown Qwen3 version $version")) end fn_to_dtype(::Type{T}) where {T} = T fn_to_dtype(::typeof(f16)) = Float16 fn_to_dtype(::typeof(f32)) = Float32 fn_to_dtype(::typeof(f64)) = Float64 fn_to_dtype(::typeof(bf16)) = BFloat16 ``` ## Model Definition {#Model-Definition} ```julia function Qwen3MLP(cfg::Qwen3Config) return Chain(; proj=Parallel( .*; gate_proj=Dense(cfg.emb_dim => cfg.hidden_dim, swish; use_bias=false), up_proj=Dense(cfg.emb_dim => cfg.hidden_dim; use_bias=false), ), down_proj=Dense(cfg.hidden_dim => cfg.emb_dim; use_bias=false), name="Qwen3MLP", ) end Qwen3RMSNorm(emb_dim::Int, eps) = AlternatePrecision{Float32}(RMSNorm(emb_dim; epsilon=eps)) @concrete struct GroupedQueryAttention <: AbstractLuxContainerLayer{( :q_proj, :k_proj, :v_proj, :o_proj, :q_norm, :k_norm )} q_proj k_proj v_proj o_proj q_norm k_norm d_in::Int num_heads::Int num_kv_groups::Int head_dim::Int end function GroupedQueryAttention(d_in, num_heads, num_kv_groups; head_dim=nothing) @assert num_heads % num_kv_groups == 0 "num_heads must be divisible by num_kv_groups" if head_dim === nothing @assert d_in % num_heads == 0 "`d_in` must be divisible by `num_heads` if \ `head_dim` is not set" head_dim = d_in ÷ num_heads end d_out = num_heads * head_dim return GroupedQueryAttention( Dense(d_in, d_out; use_bias=false), Dense(d_in, num_kv_groups * head_dim; use_bias=false), Dense(d_in, num_kv_groups * head_dim; use_bias=false), Dense(d_out, d_in; use_bias=false), Qwen3RMSNorm(head_dim, 1.0f-6), Qwen3RMSNorm(head_dim, 1.0f-6), d_in, num_heads, num_kv_groups, head_dim, ) end function apply_rope(x::AbstractArray{T}, cos_cache, sin_cache) where {T} return T.(apply_rotary_embedding(x, cos_cache, sin_cache; seq_dim=3)) end function (attn::GroupedQueryAttention)((x, cos_cache, sin_cache), ps, st::NamedTuple) _, num_tokens, B = size(x) # apply projections queries, st_q_proj = attn.q_proj(x, ps.q_proj, st.q_proj) keys, st_k_proj = attn.k_proj(x, ps.k_proj, st.k_proj) values, st_v_proj = attn.v_proj(x, ps.v_proj, st.v_proj) # reshape and permute to (head_dim, num_heads/num_kv_groups, num_tokens, batch) queries = reshape(queries, attn.head_dim, attn.num_heads, num_tokens, B) keys = reshape(keys, attn.head_dim, attn.num_kv_groups, num_tokens, B) values = reshape(values, attn.head_dim, attn.num_kv_groups, num_tokens, B) # apply normalization queries, st_q_norm = attn.q_norm(queries, ps.q_norm, st.q_norm) keys, st_k_norm = attn.k_norm(keys, ps.k_norm, st.k_norm) # apply RoPE queries = apply_rope(queries, cos_cache, sin_cache) keys = apply_rope(keys, cos_cache, sin_cache) # attention context = reshape( scaled_dot_product_attention( queries, keys, values; head_dim=1, token_dim=3, is_causal=true )[1], attn.head_dim * attn.num_heads, num_tokens, B, ) # output projection proj, st_o_proj = attn.o_proj(context, ps.o_proj, st.o_proj) return ( proj, (; q_proj=st_q_proj, k_proj=st_k_proj, v_proj=st_v_proj, o_proj=st_o_proj, q_norm=st_q_norm, k_norm=st_k_norm, ), ) end @concrete struct Qwen3Attention <: AbstractLuxContainerLayer{( :self_attn, :mlp, :input_layernorm, :post_attention_layernorm )} self_attn <: GroupedQueryAttention mlp input_layernorm post_attention_layernorm end function Qwen3Attention(cfg::Qwen3Config) return Qwen3Attention( GroupedQueryAttention(cfg.emb_dim, cfg.n_heads, cfg.n_kv_groups; cfg.head_dim), Qwen3MLP(cfg), Qwen3RMSNorm(cfg.emb_dim, 1.0f-6), Qwen3RMSNorm(cfg.emb_dim, 1.0f-6), ) end function (block::Qwen3Attention)((x, cos_cache, sin_cache), ps, st::NamedTuple) # shortcut connection for attention block shortcut = x x, st_norm1 = block.input_layernorm(x, ps.input_layernorm, st.input_layernorm) x, st_attn = block.self_attn((x, cos_cache, sin_cache), ps.self_attn, st.self_attn) x = x .+ shortcut # shortcut connection for feed-forward block shortcut = x x, st_norm2 = block.post_attention_layernorm( x, ps.post_attention_layernorm, st.post_attention_layernorm ) x, st_ff = block.mlp(x, ps.mlp, st.mlp) x = x .+ shortcut return ( x, (; self_attn=st_attn, mlp=st_ff, input_layernorm=st_norm1, post_attention_layernorm=st_norm2, ), ) end @concrete struct Qwen3 <: AbstractLuxContainerLayer{(:embed_tokens, :blocks, :norm, :lm_head)} embed_tokens blocks norm lm_head cfg::Qwen3Config end function Qwen3(cfg::Qwen3Config) return Qwen3( Embedding(cfg.vocab_size => cfg.emb_dim), Tuple([Qwen3Attention(cfg) for _ in 1:(cfg.n_layers)]), Qwen3RMSNorm(cfg.emb_dim, 1.0f-6), Dense(cfg.emb_dim, cfg.vocab_size; use_bias=false), cfg, ) end function LuxCore.initialstates(rng::AbstractRNG, m::Qwen3) head_dim = m.cfg.head_dim === nothing ? m.cfg.emb_dim ÷ m.cfg.n_heads : m.cfg.head_dim (; cos_cache, sin_cache) = compute_rotary_embedding_params( head_dim, m.cfg.context_length; base=m.cfg.rope_base, dtype=Float32 ) return (; cos_cache, sin_cache, embed_tokens=LuxCore.initialstates(rng, m.embed_tokens), blocks=LuxCore.initialstates(rng, m.blocks), norm=LuxCore.initialstates(rng, m.norm), lm_head=LuxCore.initialstates(rng, m.lm_head), ) end function (qwen3::Qwen3)(in_idx, ps, st::NamedTuple) x, st_embed_tokens = qwen3.embed_tokens(in_idx, ps.embed_tokens, st.embed_tokens) st_blocks = () for (i, block) in enumerate(qwen3.blocks) x, st_block_new = block((x, st.cos_cache, st.sin_cache), ps.blocks[i], st.blocks[i]) st_blocks = (st_blocks..., st_block_new) end x, st_norm = qwen3.norm(x, ps.norm, st.norm) logits, st_lm_head = qwen3.lm_head( fn_to_dtype(qwen3.cfg.dtype).(x), ps.lm_head, st.lm_head ) return ( logits, (; cos_cache=st.cos_cache, sin_cache=st.sin_cache, embed_tokens=st_embed_tokens, blocks=st_blocks, norm=st_norm, lm_head=st_lm_head, ), ) end ``` ## Model Weights and Tokenizer from HuggingFace {#Model-Weights-and-Tokenizer-from-HuggingFace} ```julia function download_qwen3_weights_from_huggingface(cfg::Qwen3Config) return download_qwen3_weights_from_huggingface(cfg.reasoning_model, cfg.version) end function download_qwen3_weights_from_huggingface(use_reasoning_model::Bool, version::String) repo_id = "Qwen/Qwen3-$(version)" * (use_reasoning_model ? "" : "-Base") local_dir = @get_scratch!("Qwen3-$(version)-$(use_reasoning_model)") tokenizer_file = huggingface_hub.hf_hub_download(; repo_id=repo_id, filename="tokenizer.json", local_dir=local_dir ) if version == "0.6B" weights_file = huggingface_hub.hf_hub_download(; repo_id=repo_id, filename="model.safetensors", local_dir=local_dir ) weights_dict = load_safetensors(string(weights_file)) else repo_dir = huggingface_hub.snapshot_download(; repo_id=repo_id, local_dir=local_dir) index_path = joinpath(string(repo_dir), "model.safetensors.index.json") index = JSON3.read(index_path) weights_dict = Dict() for filename in Set(values(index["weight_map"])) shard_path = joinpath(string(repo_dir), filename) shard = load_safetensors(shard_path) merge!(weights_dict, shard) end end return weights_dict, string(tokenizer_file), repo_id end ``` ## Qwen3 Tokenizer {#Qwen3-Tokenizer} ```julia struct Qwen3Tokenizer tokenizer::Tokenizer special_to_id::Dict{String,Int32} pad_token_id::Int32 eos_token_id::Int32 apply_chat_template::Bool add_generation_prompt::Bool add_thinking::Bool end function Base.show(io::IO, tokenizer::Qwen3Tokenizer) return print( io, "Qwen3Tokenizer(apply_chat_template=$(tokenizer.apply_chat_template), add_generation_prompt=$(tokenizer.add_generation_prompt), add_thinking=$(tokenizer.add_thinking))", ) end const SPECIALS = [ "<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|object_ref_start|>", "<|object_ref_end|>", "<|box_start|>", "<|box_end|>", "<|quad_start|>", "<|quad_end|>", "<|vision_start|>", "<|vision_end|>", "<|vision_pad|>", "<|image_pad|>", "<|video_pad|>", ] const SPLIT_RE = r"(<\|[^>]+?\|>)" token_to_id(tokenizer::Qwen3Tokenizer, s) = token_to_id(tokenizer.tokenizer, s) function token_to_id(tokenizer::Tokenizer, s) return pyconvert(Int32, tokenizer.py_tokenizer.token_to_id(s)) + Int32(1) end function split_with_delims(text::String, re::Regex) parts = String[] last_end = 1 for m in eachmatch(re, text) if m.offset > last_end push!(parts, text[last_end:(m.offset - 1)]) elseif m.offset == 1 push!(parts, "") end push!(parts, m.match) last_end = m.offset + length(m.match) end if last_end ≤ lastindex(text) push!(parts, text[last_end:end]) end return parts end function Qwen3Tokenizer( tokenizer_file_path::String; repo_id=nothing, apply_chat_template::Bool=true, add_generation_prompt::Bool=false, add_thinking::Bool=false, ) tok = HuggingFaceTokenizers.from_file(Tokenizer, tokenizer_file_path) special_to_id = Dict(s => token_to_id(tok, s) for s in SPECIALS) pad_token_id = special_to_id["<|endoftext|>"] eos_token_id = pad_token_id if repo_id !== nothing && !occursin("Base", repo_id) eos_token = "<|im_end|>" else eos_token = "<|endoftext|>" end if haskey(special_to_id, eos_token) eos_token_id = special_to_id[eos_token] end return Qwen3Tokenizer( tok, special_to_id, pad_token_id, eos_token_id, apply_chat_template, add_generation_prompt, add_thinking, ) end function wrap_chat(tokenizer::Qwen3Tokenizer, user_msg::AbstractString) s = "<|im_start|>user\n$(user_msg)<|im_end|>\n" if tokenizer.add_generation_prompt s *= "<|im_start|>assistant" if tokenizer.add_thinking s *= "\n" else s *= "\n\n\n\n\n" end end return s end function HuggingFaceTokenizers.encode( tok::Qwen3Tokenizer, text; chat_wrapped::Bool=tok.apply_chat_template ) stripped = strip(text) if haskey(tok.special_to_id, stripped) && !occursin('\n', stripped) return [tok.special_to_id[stripped]] end chat_wrapped && (text = wrap_chat(tok, text)) ids = Int32[] for part in filter(!isempty, split_with_delims(text, SPLIT_RE)) if haskey(tok.special_to_id, part) push!(ids, tok.special_to_id[part]) else append!(ids, encode(tok.tokenizer, string(part)).ids .+ Int16(1)) end end return ids end function HuggingFaceTokenizers.decode(tok::Qwen3Tokenizer, ids::Vector{<:Integer}) return decode(tok.tokenizer, ids .- Int16(1); skip_special_tokens=false) end ``` ## Pretrained Model Weights {#Pretrained-Model-Weights} ```julia get_weights_tensor(tensor::AbstractArray, ::Type{T}) where {T} = collect(T, tensor) function get_weights_tensor(dict, key, dtype::Type{T}, dev; permute::Bool=false) where {T} tensor = dict[key] if permute tensor = permutedims(tensor, Tuple(reverse(1:ndims(tensor)))) end return get_weights_tensor(tensor, dtype) |> dev end function load_weights_from_dict(weights_dict, cfg::Qwen3Config, dev) dtype = fn_to_dtype(cfg.dtype) function get_tensor(key; kwargs...) return get_weights_tensor(weights_dict, key, dtype, dev; kwargs...) end embed_tokens = (; weight=get_tensor("model.embed_tokens.weight"; permute=true)) blocks = Vector{Any}(undef, cfg.n_layers) for l in 1:(cfg.n_layers) prefix = "model.layers.$(l - 1)" sa_prefix = "$(prefix).self_attn" blocks[l] = (; self_attn=merge( NamedTuple( k => (; weight=get_tensor("$(sa_prefix).$k.weight")) for k in (:q_proj, :k_proj, :v_proj, :o_proj) ), NamedTuple( k => (; scale=get_tensor("$(sa_prefix).$k.weight")) for k in (:q_norm, :k_norm) ), ), mlp=(; proj=(; gate_proj=(; weight=get_tensor("$(prefix).mlp.gate_proj.weight")), up_proj=(; weight=get_tensor("$(prefix).mlp.up_proj.weight")), ), down_proj=(; weight=get_tensor("$(prefix).mlp.down_proj.weight")), ), input_layernorm=(; scale=get_tensor("$(prefix).input_layernorm.weight")), post_attention_layernorm=(; scale=get_tensor("$(prefix).post_attention_layernorm.weight") ), ) end blocks = Tuple(blocks) norm = (; scale=get_weights_tensor(weights_dict, "model.norm.weight", dtype, dev)) if haskey(weights_dict, "lm_head.weight") lm_head = (; weight=get_weights_tensor(weights_dict, "lm_head.weight", dtype, dev)) else # Weight tying with the embedding matrix. We will share the weights here to # reduce memory usage. lm_head = (; weight=transpose(embed_tokens.weight)) end return (; embed_tokens, blocks, norm, lm_head) end function setup_model( version::String, dev; weights_dict::Union{Nothing,Dict}=nothing, kwargs... ) return setup_model(Qwen3Config(version; kwargs...), dev; weights_dict) end function setup_model(cfg::Qwen3Config, dev; weights_dict::Union{Nothing,Dict}=nothing) model = Qwen3(cfg) st = Lux.initialstates(Random.default_rng(), model) |> dev if weights_dict !== nothing ps = load_weights_from_dict(weights_dict, cfg, dev) else ps = Lux.initialparameters(Random.default_rng(), model) |> dev ps = ps |> cfg.dtype end return model, ps, st end ``` ## Running the model without dynamic sizes {#Running-the-model-without-dynamic-sizes} ```julia function get_padded_size(seq_len::Int, context_length::Int) return min(max(512, nextpow(2, seq_len)), context_length) end function padded_input_and_mask_len(x::AbstractMatrix, v, cfg::Qwen3Config, pad_token_id) return padded_input_and_mask_len( x, v, get_padded_size(size(x, 1) + v !== nothing, cfg.context_length), pad_token_id ) end function padded_input_and_mask_len(x::AbstractMatrix, v, padded_sz::Int, pad_token_id) if padded_sz > size(x, 1) x_padded = similar(x, (padded_sz, size(x, 2))) x_padded[1:size(x, 1), :] .= x if v === nothing x_padded[(size(x, 1) + 1):end, :] .= pad_token_id else x_padded[(size(x, 1) + 1), :] = v[1, :] x_padded[(size(x, 1) + 2):end, :] .= pad_token_id end else x_padded = x end return ( x_padded, Reactant.promote_to( Reactant.TracedRNumber{Int32}, padded_sz - (size(x, 1) + (v !== nothing)) ), ) end ``` ## Helpers to generate text {#Helpers-to-generate-text} ```julia function predict_next_token( model, token_ids::AbstractMatrix{T}, input_mask_len, ps, st ) where {T} logits, stₙ = model(token_ids, ps, st) predictions = T.(argmax(logits[:, end - input_mask_len, :]; dims=1)) predictions = mod1.(predictions, T(size(logits, 1))) return predictions, stₙ end function update_token_ids_and_mask!( padded_token_ids, input_mask_len, cur_num_tokens, next_token ) next_token_idx = safe_increment(cur_num_tokens) padded_token_ids[next_token_idx, :] = next_token[1, :] return input_mask_len - eltype(input_mask_len)(1), next_token_idx end function update_token_ids_with_shift!(token_ids, next_token) token_ids[1:(end - 1), :] = token_ids[2:end, :] token_ids[end, :] = next_token[1, :] return nothing end safe_increment(x) = x + one(x) mutable struct CachedReactantThunks cache::Dict{Qwen3Config,Dict{Int,NTuple{3,Reactant.Compiler.Thunk}}} increment_fn::Union{Nothing,Reactant.Compiler.Thunk} end function CachedReactantThunks() return CachedReactantThunks( Dict{Qwen3Config,Dict{Int,NTuple{3,Reactant.Compiler.Thunk}}}(), nothing ) end function cache_and_retrieve!( cache::CachedReactantThunks, len::Integer, model::Qwen3, padded_token_ids, input_mask_len, ps, st, next_token, cur_num_tokens_traced, ) if haskey(cache.cache, model.cfg) && haskey(cache.cache[model.cfg], len) return cache.cache[model.cfg][len] end println() @warn "Compiling Qwen3 generation loop for $(model.cfg.version) with $(len) tokens. \ This might take a while... (However this is only done once per model per length)" predict_next_token_compiled = @compile predict_next_token( model, padded_token_ids, input_mask_len, ps, st ) update_fn1! = @compile update_token_ids_and_mask!( padded_token_ids, input_mask_len, cur_num_tokens_traced, next_token ) update_fn2! = @compile update_token_ids_with_shift!(padded_token_ids, next_token) if !haskey(cache.cache, model.cfg) cache.cache[model.cfg] = Dict{Int,NTuple{3,Reactant.Compiler.Thunk}}() end return cache.cache[model.cfg][len] = ( predict_next_token_compiled, update_fn1!, update_fn2! ) end const CACHED_THUNKS = CachedReactantThunks() generate_text(args...; kwargs...) = generate_text!(CACHED_THUNKS, args...; kwargs...) function generate_text!( compile_cache::CachedReactantThunks, model::Qwen3, prompt::String, ps, st, max_new_tokens, tokenizer, ) token_ids = Reactant.to_rarray(reshape(encode(tokenizer, prompt), :, 1)) # TODO: compile the generation loop with Reactant # TODO: implement some simple KV caching cur_num_tokens = size(token_ids, 1) max_context_length = model.cfg.context_length cur_compiled_fn_token_len = get_padded_size(cur_num_tokens, max_context_length) padded_token_ids, input_mask_len = @jit padded_input_and_mask_len( token_ids, nothing, cur_compiled_fn_token_len, tokenizer.pad_token_id ) cur_num_tokens_traced = ConcreteRNumber{Int32}(cur_num_tokens) next_token = get_device(ps)(rand(Int32, 1, size(padded_token_ids, 2))) (predict_next_token_compiled, update_fn1!, update_fn2!) = cache_and_retrieve!( compile_cache, cur_compiled_fn_token_len, model, padded_token_ids, input_mask_len, ps, st, next_token, cur_num_tokens_traced, ) if compile_cache.increment_fn === nothing compile_cache.increment_fn = @compile safe_increment(cur_num_tokens_traced) end start_time = time() compile_time = 0.0 ntokens_generated = 0 for _ in 1:max_new_tokens new_compiled_fn_token_len = get_padded_size(cur_num_tokens, max_context_length) if new_compiled_fn_token_len != cur_compiled_fn_token_len compile_start_time = time() cur_compiled_fn_token_len = new_compiled_fn_token_len padded_token_ids, input_mask_len = @jit padded_input_and_mask_len( padded_token_ids, next_token, cur_compiled_fn_token_len, tokenizer.pad_token_id, ) (predict_next_token_compiled, update_fn1!, update_fn2!) = cache_and_retrieve!( compile_cache, cur_compiled_fn_token_len, model, padded_token_ids, input_mask_len, ps, st, next_token, cur_num_tokens_traced, ) compile_time += time() - compile_start_time end next_token, st = predict_next_token_compiled( model, padded_token_ids, input_mask_len, ps, st ) ntokens_generated += 1 next_token_jl = vec(Array(next_token)) if tokenizer.eos_token_id !== nothing && all(next_token_jl .== tokenizer.eos_token_id) break end print(decode(tokenizer, next_token_jl)) if cur_num_tokens >= max_context_length update_fn2!(padded_token_ids, next_token) elseif new_compiled_fn_token_len > cur_num_tokens input_mask_len, cur_num_tokens_traced = update_fn1!( padded_token_ids, input_mask_len, cur_num_tokens_traced, next_token ) else cur_num_tokens_traced = compile_cache.increment_fn(cur_num_tokens_traced) end cur_num_tokens += 1 end total_time = time() - start_time println() return ntokens_generated / (total_time - compile_time) end ``` ## Entry Point {#Entry-Point} ```julia function run_model_selection() printstyled("Which model do you want to run? \n"; color=:cyan, bold=true) choices = ["0.6B", "1.7B", "4B", "8B", "14B", "32B"] for (i, choice) in enumerate(choices) printstyled(" $(i). $(choice)\n"; color=:light_blue) end printstyled(" Enter your choice: "; color=:cyan) choice = parse(Int, readline(stdin)) if choice ∉ 1:length(choices) error("Invalid choice: $(choice). Expected an integer between 1 and \ $(length(choices))") end printstyled("Do you want to use the reasoning model? [y/N] "; color=:cyan) reasoning = readline(stdin) == "y" println() return choices[choice], reasoning end function get_model_and_tokenizer(version, reasoning) cfg = Qwen3Config(version; reasoning_model=reasoning) rdev = reactant_device(; force=true) weights_dict, tokenizer_file, repo_id = download_qwen3_weights_from_huggingface(cfg) tokenizer = Qwen3Tokenizer( tokenizer_file; repo_id, add_generation_prompt=cfg.reasoning_model, add_thinking=cfg.reasoning_model, ) model, ps, st = setup_model(cfg, rdev; weights_dict) return model, ps, st, tokenizer end function main() @info "Text Generation with Qwen-3 powered by Lux, Reactant & XLA." version, reasoning = run_model_selection() model, ps, st, tokenizer = get_model_and_tokenizer(version, reasoning) while true printstyled( "Prompt (type \"exit\" to quit the program or \ \"model selection\" to change the model): "; color=:cyan, bold=true, ) prompt = readline(stdin) prompt == "exit" && break if prompt == "model selection" version, reasoning = run_model_selection() model, ps, st, tokenizer = get_model_and_tokenizer(version, reasoning) continue end tokens_per_second = generate_text(model, prompt, ps, st, 100_000, tokenizer) println("\nTokens per second: $tokens_per_second\n\n") end return nothing end if abspath(PROGRAM_FILE) == @__FILE__ main() end ``` *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/manual/interface.md --- # Lux Interface {#lux-interface} ::: tip Lux.jl vs LuxCore.jl If you just want to define compatibility with Lux without actually using any of the other functionality provided by Lux (like layers), it is recommended to depend on `LuxCore.jl` instead of `Lux.jl`. `LuxCore.jl` is a significantly lighter dependency. ::: Following this interface provides the ability for frameworks built on top of Lux to be cross compatible. Additionally, any new functionality built into Lux, will just work for your framework. ::: tip `@compact` macro While writing out a custom struct and defining dispatches manually is a good way to understand the interface, it is not the most concise way. We recommend using the [`Lux.@compact`](/api/Lux/utilities#Lux.@compact) macro to define layers which makes handling the states and parameters downright trivial. ::: ## Layer Interface {#Layer-Interface} ### Singular Layer {#Singular-Layer} If the layer doesn't contain any other Lux layer, then it is a `Singular Layer`. This means it should optionally subtype `Lux.AbstractLuxLayer` but mandatorily define all the necessary functions mentioned in the docstrings. Consider a simplified version of [`Dense`](/api/Lux/layers#Lux.Dense) called `Linear`. First, setup the architectural details for this layer. Note, that the architecture doesn't contain any mutable structure like arrays. When in doubt, remember, once constructed a model architecture cannot change. ::: tip Tip For people coming from Flux.jl background, this might be weird. We recommend checking out [the Flux to Lux migration guide](/manual/migrate_from_flux#migrate-from-flux) first before proceeding. ::: ```julia using LuxCore, Random, WeightInitializers # Importing `Lux` also gives you access to `LuxCore` struct Linear{F1, F2} <: LuxCore.AbstractLuxLayer in_dims::Int out_dims::Int init_weight::F1 init_bias::F2 end function Linear(in_dims::Int, out_dims::Int; init_weight=glorot_uniform, init_bias=zeros32) return Linear{typeof(init_weight), typeof(init_bias)}(in_dims, out_dims, init_weight, init_bias) end l = Linear(2, 4) ``` ```ansi Main.Linear{typeof(glorot_uniform), typeof(zeros32)}(2, 4, WeightInitializers.glorot_uniform, WeightInitializers.zeros32) ``` Next, we need to implement functions which return the parameters and states for the layer. In case of `Linear`, the parameters are `weight` and `bias` while the states are empty. States become important when defining layers like [`BatchNorm`](/api/Lux/layers#Lux.BatchNorm), [`WeightNorm`](/api/Lux/layers#Lux.WeightNorm), etc. The recommended data structure for returning parameters is a NamedTuple, though anything satisfying the [Parameter Interface](/manual/interface#parameter-interface) is valid. ```julia function LuxCore.initialparameters(rng::AbstractRNG, l::Linear) return (weight=l.init_weight(rng, l.out_dims, l.in_dims), bias=l.init_bias(rng, l.out_dims, 1)) end LuxCore.initialstates(::AbstractRNG, ::Linear) = NamedTuple() ``` You could also implement `LuxCore.parameterlength` and `LuxCore.statelength` to prevent wasteful reconstruction of the parameters and states. ```julia # This works println("Parameter Length: ", LuxCore.parameterlength(l), "; State Length: ", LuxCore.statelength(l)) # But still recommended to define these LuxCore.parameterlength(l::Linear) = l.out_dims * l.in_dims + l.out_dims LuxCore.statelength(::Linear) = 0 ``` ```ansi Parameter Length: 12; State Length: 0 ``` ::: tip No RNG in `parameterlength` and `statelength` You might notice that we don't pass in a `RNG` for these functions. If your parameter length and/or state length depend on a random number generator, you should think **really hard** about what you are trying to do and why. ::: Now, we need to define how the layer works. For this you make your layer a function with exactly 3 arguments – `x` the input, `ps` the parameters, and `st` the states. This function must return two things – `y` the output, and `st_new` the updated state. ```julia function (l::Linear)(x::AbstractMatrix, ps, st::NamedTuple) y = ps.weight * x .+ ps.bias return y, st end ``` Finally, let's run this layer. If you have made this far into the documentation, we don't feel you need a refresher on that. But if you do see the [Quickstart](/introduction/index#Quickstart). ```julia rng = Random.default_rng() Random.seed!(rng, 0) ps, st = LuxCore.setup(rng, l) println("Parameter Length: ", LuxCore.parameterlength(l), "; State Length: ", LuxCore.statelength(l)) x = randn(rng, Float32, 2, 1) LuxCore.apply(l, x, ps, st) # or `l(x, ps, st)` ``` ```ansi (Float32[-0.15276335; 0.45325348; 1.0207279; 0.78226817;;], NamedTuple()) ``` ### Container Layer {#Container-Layer} If your layer comprises of other Lux layers, then it is a `Container Layer`. Note that you could treat it as a [`Singular Layer`](/manual/interface#singular-layer), and it is still fine. FWIW, if you cannot subtype your layer with `LuxCore.AbstractLuxContainerLayer` then you should go down the [`Singular Layer`](/manual/interface#singular-layer) route. But subtyping allows us to bypass some of these common definitions. Let us now define a layer, which is basically a composition of two linear layers. ::: tip Wrapper Layer If you are defining a layer that is a wrapper around another layer, then you should subtype `LuxCore.AbstractLuxWrapperLayer` instead of `LuxCore.AbstractLuxContainerLayer`. The only difference from a container layer is that it can wrap a single layer and the parameter/state structure is exactly the same as the wrapped layer. ::: ```julia struct ComposedLinear{L1, L2} <: LuxCore.AbstractLuxContainerLayer{(:linear_1, :linear_2)} linear_1::L1 linear_2::L2 end function (cl::ComposedLinear)(x::AbstractMatrix, ps, st::NamedTuple) # To access the parameters and states for `linear_1` we do `ps.linear_1` and # `st.linear_1`. Similarly for `linear_2` y, st_l1 = cl.linear_1(x, ps.linear_1, st.linear_1) y, st_l2 = cl.linear_2(y, ps.linear_2, st.linear_2) # Finally, we need to return the new state which has the exact structure as `st` return y, (linear_1 = st_l1, linear_2 = st_l2) end ``` Here, you will notice we have passed `(:linear_1, :linear_2)` to the supertype. It essentially informs the type that, `.linear_1` and `.linear_2` are Lux layers and we need to construct parameters and states for those. Let's construct these and see: ```julia model = ComposedLinear(Linear(2, 4), Linear(4, 2)) display(model) ps, st = LuxCore.setup(rng, model) println("Parameters: ", ps) println("States: ", st) println("Parameter Length: ", LuxCore.parameterlength(model), "; State Length: ", LuxCore.statelength(model)) x = randn(rng, Float32, 2, 1) LuxCore.apply(model, x, ps, st) # or `model(x, ps, st)` ``` ```ansi (Float32[1.3410565; 0.78000563;;], (linear_1 = NamedTuple(), linear_2 = NamedTuple())) ``` ## Parameter Interface {#Parameter-Interface} We accept any parameter type as long as we can fetch the parameters using `getproperty(obj, :parameter_name)`. This allows us to simultaneously support `NamedTuple`s and `ComponentArray`s. Let us go through a concrete example of what it means. Consider [`Dense`](/api/Lux/layers#Lux.Dense) which expects two parameters named `weight` and `bias`. ::: tip Automatic Differentiation If you are defining your own parameter type, it is your responsibility to make sure that it works with the AutoDiff System you are using. ::: ```julia using Lux, Random d = Dense(2, 3) rng = Random.default_rng() Random.seed!(rng, 0) ps_default, st = LuxCore.setup(rng, d) x = randn(rng, Float32, 2, 1) println("Result with `NamedTuple` parameters: ", first(d(x, ps_default, st))) ``` ```ansi Result with `NamedTuple` parameters: Float32[-0.08713347; -0.4851346; -0.8490221;;] ``` Let, us define a custom parameter type with fields `myweight` and `mybias` but if we try to access `weight` we get back `myweight`, similar for `bias`. ::: warning Beware! This is for demonstrative purposes, don't try this at home! ::: ```julia struct DenseLayerParameters{W, B} myweight::W mybias::B end function Base.getproperty(ps::DenseLayerParameters, x::Symbol) if x == :weight return getfield(ps, :myweight) elseif x == :bias return getfield(ps, :mybias) end return getfield(ps, x) end ps = DenseLayerParameters(ps_default.weight, ps_default.bias) println("Result with `DenseLayerParameters` parameters: ", first(d(x, ps, st))) ``` ```ansi Result with `DenseLayerParameters` parameters: Float32[0.23710957; 0.1003911; -0.57671577;;] ``` The takeaway from this shouldn't be – *lets define weird parameter types*. Simply because you can do weird things like this doesn't mean you should, since it only leads to bugs. Instead this shows the flexibility you have for how your parameters can be structured. ## State Interface {#State-Interface} States are always type constrained to be `NamedTuple`. The structure of the input state **must** match that of the output state, i.e. `keys(st_in) == keys(st_out)`. This doesn't imply that types of the input and output state match. To generate efficient code, we often do dispatch on the state, for example, [`Lux.Dropout`](/api/Lux/layers#Lux.Dropout), [`BatchNorm`](/api/Lux/layers#Lux.BatchNorm), etc. --- --- url: /dev/manual/freezing_model_parameters.md --- # Freezing Model Parameters {#freezing-model-parameters} ::: warning Warning API for freezing parameters should be considered experimental at this point. ::: In this manual entry, we will go over how to freeze certain parameters in a model. ## Freezing Layers of a Particular Kind {#Freezing-Layers-of-a-Particular-Kind} To freeze a particular kind of layer, let's say [`Dense`](/api/Lux/layers#Lux.Dense) in the following example. We can use [`Lux.Experimental.layer_map`](/api/Lux/contrib#Lux.Experimental.layer_map) and freeze layers if they are of type `Dense`. ```julia using Lux, Random, Functors rng = Xoshiro(0) model = Chain(Dense(3, 4), Chain(Dense(4, 4), Dropout(0.5f0), BatchNorm(4)), Dense(4, 1)) ps, st = Lux.setup(rng, model) x = randn(rng, Float32, 3, 2) model(x, ps, st) function freeze_dense(d::Lux.Dense, ps, st, path) return Lux.Experimental.freeze(d, ps, st, (:weight, :bias)) end freeze_dense(l, ps, st, path) = (l, ps, st) model_frozen, ps_frozen, st_frozen = Lux.Experimental.layer_map(freeze_dense, model, ps, st) model_frozen(x, ps_frozen, st_frozen) ``` ```ansi (Float32[0.6886741 -1.2361472], (layer_1 = (frozen_params = (weight = Float32[-0.028461456 -0.5999714 -0.3850993; -0.18860114 0.72428167 0.32322538; -0.965117 -0.4585489 -0.32623518; -0.86290836 -0.82805836 -0.7673453], bias = Float32[0.4216236, -0.4510427, -0.097253, 0.23325463]), states = NamedTuple()), layer_2 = (layer_1 = (frozen_params = (weight = Float32[-0.680748 0.1764085 0.34383082 0.6469914; -0.13819042 -0.109261915 -0.6143286 -0.21672015; -0.20881107 0.70390546 0.48137343 0.25662464; 0.38187847 0.05779423 -0.35181466 -0.096988946], bias = Float32[0.41246277, 0.4318977, -0.4305781, 0.3367505]), states = NamedTuple()), layer_2 = (rng = Random.Xoshiro(0x4fa3403dd074e603, 0x12c522b8034ae186, 0x8e0c3a65079041bb, 0x21617f7747d97206, 0x22a21880af5dc689), training = Val{true}()), layer_3 = (running_mean = Float32[0.01965834, 0.0, 0.0, 0.015937408], running_var = Float32[0.90772897, 0.9, 0.9, 0.90508], training = Val{true}())), layer_3 = (frozen_params = (weight = Float32[0.7794657 0.8337032 0.6323408 -0.18308182], bias = Float32[-0.27373654]), states = NamedTuple()))) ``` ## Freezing by Layer Name {#Freezing-by-Layer-Name} When the function in `layer_map` is called, the 4th argument is the name of the layer. For example, if you want to freeze the 1st layer inside the inner Chain. The name for this would be `layer_2.layer_1`. :::code-group ```julia [Freezing by Layer Name] function freeze_by_name(d, ps, st, name::KeyPath) name == KeyPath(:layer_2, :layer_1) && return Lux.Experimental.freeze(d, ps, st, (:weight, :bias)) return d, ps, st end ``` ```julia [Freezing by Layer Type] function freeze_dense(d::Dense, ps, st, _) return Lux.Experimental.freeze(d, ps, st, (:weight, :bias)) end freeze_dense(l, ps, st, _) = (l, ps, st) ``` ::: ## Freezing Part of the Parameters {#Freezing-Part-of-the-Parameters} Instead of freezing all the parameters, we can simply specify `(:weight,)` to freeze only the `weight` parameter while training the `bias` parameter. ::: code-group ```julia [Freezing Some Parameters of a Layer] function freeze_by_name(d, ps, st, name::KeyPath) name == KeyPath(:layer_2, :layer_1) && return Lux.Experimental.freeze(d, ps, st, (:weight,)) return d, ps, st end ``` ```julia [Freezing All Parameters of a Layer] function freeze_by_name(d, ps, st, name::KeyPath) name == KeyPath(:layer_2, :layer_1) && return Lux.Experimental.freeze(d, ps, st, (:weight, :bias)) return d, ps, st end ``` ::: ## Freezing Part of a Chain {#Freezing-Part-of-a-Chain} ```julia using Lux, Random rng = Random.default_rng() Random.seed!(rng, 0) model = Chain(Dense(3, 4), Dense(4, 4), Dropout(0.5f0), BatchNorm(4), Dense(4, 1)) model_frozen = Chain(model[1:2], Lux.Experimental.freeze(model[3:4]), model[5]) ps, st = Lux.setup(rng, model_frozen) x = randn(rng, Float32, 3, 2) model_frozen(x, ps, st) ``` ```ansi (Float32[0.7429947 -1.2904677], (layer_1 = (layer_1 = NamedTuple(), layer_2 = NamedTuple()), layer_2 = (frozen_params = (layer_3 = NamedTuple(), layer_4 = (scale = Float32[1.0, 1.0, 1.0, 1.0], bias = Float32[0.0, 0.0, 0.0, 0.0])), states = (layer_3 = (rng = Random.TaskLocalRNG(), training = Val{true}()), layer_4 = (running_mean = Float32[0.0, 0.048522998, 0.0, 0.015937408], running_var = Float32[0.9, 0.9470896, 0.9, 0.90508], training = Val{true}()))), layer_3 = NamedTuple())) ``` --- --- url: /dev/manual/gpu_management.md --- # GPU Management {#GPU-Management} `Lux.jl` can handle multiple GPU backends. Currently, the following backends are supported: ```julia # Important to load trigger packages using Lux #, LuxCUDA, AMDGPU, Metal, oneAPI supported_gpu_backends() ``` ```ansi ("CUDA", "AMDGPU", "Metal", "oneAPI", "OpenCL") ``` ::: tip GPU Support via Reactant If you are using Reactant, you can use the [`reactant_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.reactant_device) function to automatically select Reactant backend if available. Additionally to force Reactant to use `gpu`, you can run `Reactant.set_default_backend("gpu")` (this is automatic). ::: ::: danger AMD GPU Support For AMD GPUs, we **strongly recommend using Reactant** instead of native AMDGPU.jl. Native AMDGPU.jl support is experimental with known limitations including deadlocks in distributed training. Use `reactant_device()` with Reactant for better AMD GPU support. ::: ::: danger Metal Support Support for Metal GPUs should be considered extremely experimental at this point. ::: ## Automatic Backend Management (Recommended Approach) {#Automatic-Backend-Management-Recommended-Approach} Automatic Backend Management is done by two simple functions: `cpu_device` and `gpu_device`. * [`cpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.cpu_device): This is a simple function and just returns a `CPUDevice` object. ```julia cdev = cpu_device() x_cpu = randn(Float32, 3, 2) ``` ```ansi 3×2 Matrix{Float32}: 1.349 0.687216 -1.06669 0.196703 -0.00973899 0.261273 ``` * [`gpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_device): This function performs automatic GPU device selection and returns an object. 1. If no GPU is available, it returns a `CPUDevice` object. 2. If a LocalPreferences file is present, then the backend specified in the file is used. To set a backend, use `Lux.gpu_backend!()`. (a) If the trigger package corresponding to the device is not loaded, then a warning is displayed. (b) If no LocalPreferences file is present, then the first working GPU with loaded trigger package is used. ```julia gdev = gpu_device() x_gpu = x_cpu |> gdev ``` ```ansi 3×2 Matrix{Float32}: 1.349 0.687216 -1.06669 0.196703 -0.00973899 0.261273 ``` ```julia (x_gpu |> cdev) ≈ x_cpu ``` ```ansi true ``` ## Manual Backend Management {#Manual-Backend-Management} Automatic Device Selection can be circumvented by directly using `CPUDevice` and `AbstractGPUDevice` objects. ```julia cdev = cpu_device() x_cpu = randn(Float32, 3, 2) if MLDataDevices.functional(CUDADevice) gdev = CUDADevice() x_gpu = x_cpu |> gdev elseif MLDataDevices.functional(AMDGPUDevice) gdev = AMDGPUDevice() x_gpu = x_cpu |> gdev else @info "No GPU is available. Using CPU." x_gpu = x_cpu end (x_gpu |> cdev) ≈ x_cpu ``` ```ansi true ``` --- --- url: /dev/manual/weight_initializers.md --- # Initializing Weights {#Initializing-Weights} `WeightInitializers.jl` provides common weight initialization schemes for deep learning models. ```julia using WeightInitializers, Random # Fixing rng rng = Random.MersenneTwister(42) ``` ```ansi Random.MersenneTwister(42) ``` ```julia # Explicit rng call weights = kaiming_normal(rng, 2, 5) ``` ```ansi 2×5 Matrix{Float32}: 0.76545 0.255203 -0.0424012 0.643172 -0.360745 -0.0499631 0.183381 0.388315 -0.0340666 -0.54248 ``` ```julia # Default rng call weights = kaiming_normal(2, 5) ``` ```ansi 2×5 Matrix{Float32}: -0.227513 -0.265372 0.265788 1.29955 -0.192836 0.687611 0.454679 -0.433656 0.20548 0.292002 ``` ```julia # Passing kwargs (if needed) with explicit rng call weights_cl = kaiming_normal(rng; gain=1.0) weights = weights_cl(2, 5) ``` ```ansi 2×5 Matrix{Float32}: -0.094564 0.196581 0.0791126 -0.794864 0.631217 -0.381774 -0.588045 -0.113952 -0.567746 0.261636 ``` ```julia # Passing kwargs (if needed) with default rng call weights_cl = kaiming_normal(; gain=1.0) weights = weights_cl(2, 5) ``` ```ansi 2×5 Matrix{Float32}: -0.160876 -0.187646 0.18794 0.918918 -0.136356 0.486214 0.321506 -0.306641 0.145296 0.206476 ``` To generate weights directly on GPU, pass in a `CUDA.RNG`. For a complete list of supported RNG types, see [Supported RNG Types](/api/Building_Blocks/WeightInitializers#Supported-RNG-Types-WeightInit). ```julia using LuxCUDA weights = kaiming_normal(CUDA.default_rng(), 2, 5) ``` You can also generate Complex Numbers: ```julia weights = kaiming_normal(CUDA.default_rng(), ComplexF32, 2, 5) ``` ## Quick examples {#Quick-examples} The package is meant to be working with deep learning libraries such as (F)Lux. All the methods take as input the chosen `rng` type and the dimension for the array. ```julia weights = init(rng, dims...) ``` The `rng` is optional, if not specified a default one will be used. ```julia weights = init(dims...) ``` If there is the need to use keyword arguments the methods can be called with just the `rng` (optionally) and the keywords to get in return a function behaving like the two examples above. ```julia weights_init = init(rng; kwargs...) weights = weights_init(rng, dims...) # Or weights_init = init(; kwargs...) weights = weights_init(dims...) ``` --- --- url: /dev/manual/visualize_lux_models.md --- # Visualizing Lux Models using Model Explorer {#Visualizing-Lux-Models-using-Model-Explorer} We can use [model explorer](https://ai.google.dev/edge/model-explorer) to visualize both Lux models and the corresponding gradient expressions. To do this we just need to compile our model [using Reactant](/manual/compiling_lux_models#reactant-compilation) and save the resulting `mlir` file. ```julia using Lux, Reactant, Enzyme, Random dev = reactant_device(; force=true) model = Chain( Chain( Conv((3, 3), 3 => 32, relu; pad=SamePad()), BatchNorm(32), ), FlattenLayer(), Dense(32 * 32 * 32 => 32, tanh), BatchNorm(32), Dense(32 => 10) ) ps, st = Lux.setup(Random.default_rng(), model) |> dev x = randn(Float32, 32, 32, 3, 4) |> dev ``` Following instructions from [exporting lux models to stablehlo](/manual/exporting_to_jax#exporting_to_stablehlo) we can save the `mlir` file. ```julia hlo = @code_hlo model(x, ps, Lux.testmode(st)) write("exported_lux_model.mlir", string(hlo)) ``` ![](../public/model_explorer_graph_forward_pass.png) We can also visualize the gradients of the model using the same method. ```julia function ∇sumabs2_enzyme(model, x, ps, st) return Enzyme.gradient(Enzyme.Reverse, sum ∘ first ∘ Lux.apply, Const(model), x, ps, Const(st)) end hlo = @code_hlo ∇sumabs2_enzyme(model, x, ps, st) write("exported_lux_model_gradients.mlir", string(hlo)) ``` This is going to be hard to read, but you get the idea. ![](../public/model_explorer_graph_backward_pass.png) --- --- url: /dev/manual/compiling_lux_models.md --- # Compiling Lux Models using `Reactant.jl` {#reactant-compilation} Quoting the Reactant.jl Readme: > Reactant takes Julia function and compile it into MLIR and run fancy optimizations on top of it, including using EnzymeMLIR for automatic differentiation, and create relevant executables for CPU/GPU/TPU via XLA. It presently operates as a tracing system. Compiled functions will assume the same control flow pattern as was original taken by objects used at compile time, and control flow (e.g. if, for) as well as any type instabilities will be removed. The benefits of this approach is immediately making all such code available for advanced optimization with little developer effort. ```julia 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`](/api/Lux/utilities#Lux.Training.TrainState) API, skip to the [bottom of this page](/manual/compiling_lux_models#compile_lux_model_trainstate) to see how to train the model without any of this boilerplate. ::: We start by defining a simple MLP model: ```julia model = Chain( Dense(2 => 32, gelu), Dense(32 => 32, gelu), Dense(32 => 2) ) ps, st = Lux.setup(Random.default_rng(), model) ``` ```ansi ((layer_1 = (weight = Float32[0.9670442 -0.36027783; 0.078672916 0.92788666; … ; -0.65058047 -0.47006413; -0.48801818 -0.6615898], bias = Float32[-0.28780195, -0.23392133, 0.084573634, -0.59277534, -0.6795253, 0.47792822, -0.64850235, -0.55131584, -0.33091125, 0.47174177 … 0.07477753, -0.10521463, -0.45745936, 0.19031122, 0.41613227, 0.47329637, -0.68522483, -0.2834571, 0.0235815, 0.61977077]), layer_2 = (weight = Float32[-0.057887085 -0.14646342 … 0.1019723 0.14663221; 0.10022328 -0.09659223 … 0.25911948 -0.008825431; … ; -0.014519578 -0.01100632 … -0.30112675 -0.17886546; 0.21983564 -0.026677115 … -0.030971587 -0.28283697], bias = Float32[0.095548995, 0.10995198, 0.12209795, -0.14433007, 0.11754602, -0.152131, -0.10584956, 0.09469124, 0.09255884, 0.10044085 … 0.07444663, 0.11096934, 0.13462374, 0.15048876, 0.061646424, 0.004753132, 0.08162795, -0.15708117, 0.029835312, 0.005353872]), layer_3 = (weight = Float32[0.005372945 -0.18356045 … 0.052086722 0.07186686; 0.0067291846 0.020219602 … 0.0688707 -0.1961357], bias = Float32[-0.03542879, -0.041368797])), (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple())) ``` We then create a random input and output data: ```julia x = randn(Float32, 2, 32) y = x .^ 2 ``` We will use [`reactant_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.reactant_device) similar to [`gpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_device) to move the arrays to `Reactant`. ```julia const xdev = reactant_device() x_ra = x |> xdev y_ra = y |> xdev ps_ra = ps |> xdev st_ra = st |> xdev ``` First let's run the model as we would normally: ```julia pred_lux, _ = model(x, ps, Lux.testmode(st)) ``` ```ansi (Float32[0.01586983 0.010564316 … -0.4137662 0.018748946; 0.078654006 0.06953075 … -0.2340262 0.21624328], (layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple())) ``` To run it using `XLA` we need to compile the model. We can do this using the `Reactant.@compile` macro. Note that the inputs need to be moved to the device using [`reactant_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.reactant_device) first. ```julia model_compiled = @compile model(x_ra, ps_ra, Lux.testmode(st_ra)) ``` ```ansi Reactant compiled function Chain{@NamedTuple{layer_1::Dense{typeof(gelu_tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(gelu_tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(2 => 32, gelu_tanh), layer_2 = Dense(32 => 32, gelu_tanh), layer_3 = Dense(32 => 2)), nothing) (with tag ##Chain{@NamedTuple{layer_1::Dense{typeof(gelu_tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_2::Dense{typeof(gelu_tanh), Int64, Int64, Nothing, Nothing, Static.True}, layer_3::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, Nothing}((layer_1 = Dense(2 => 32, gelu_tanh), layer_2 = Dense(32 => 32, gelu_tanh), layer_3 = Dense(32 => 2)), nothing)_reactant#912992) ``` Now we can test the difference between the results: ```julia pred_compiled, _ = model_compiled(x_ra, ps_ra, Lux.testmode(st_ra)) pred_lux .- Array(pred_compiled) ``` ```ansi 2×32 Matrix{Float32}: 3.72529f-8 7.45058f-9 1.11759f-8 5.96046f-8 … 8.9407f-8 1.86265f-8 7.45058f-9 2.23517f-8 0.0 2.23517f-7 1.49012f-8 -2.98023f-8 ``` The difference is very small as we would expect. Now, let's try to differentiate the output of the model. We need to use `Enzyme.jl` to do this. ```julia function loss_function(model, ps, st, x, y) pred, _ = model(x, ps, st) return MSELoss()(pred, y) end ``` We will use `Zygote.jl` to compute the gradient of the loss function for the vanilla model. ```julia loss_function(model, ps, st, x, y) ∂ps_zyg = only(Zygote.gradient(ps -> loss_function(model, ps, st, x, y), ps)) ``` ```ansi (layer_1 = (weight = Float32[0.26016673 -0.09287718; -0.02802991 -0.013659927; … ; -0.07384164 -0.06003239; 0.042984415 0.051605415], bias = Float32[0.12923914, -0.009405827, -0.0263628, -0.014524895, 0.013915386, 0.093436174, 0.08193636, 0.0077627874, 0.001044218, 0.018755747 … 0.067041054, -0.043209504, 0.10486872, 0.014353438, 0.024228808, -0.06582927, 0.010303013, 0.098782696, 0.06784941, -0.08268406]), layer_2 = (weight = Float32[-0.004457889 0.00021804236 … -0.003327486 -0.008014375; 0.13051204 -0.0046890336 … 0.038353663 0.093302496; … ; -0.041253403 0.000887985 … 0.000747012 -0.020347353; 0.06339173 -0.0021197682 … 0.012145687 0.06877027], bias = Float32[-0.0062405514, 0.13950492, -0.22439212, -0.113269635, -0.023160838, 0.14702773, 0.03519612, 0.13981938, -0.23715457, 0.32662556 … -0.014224289, 0.009401775, 0.18295962, 0.13164552, 0.16955197, -0.110567965, -0.0074348953, 0.118868664, -0.026588853, 0.031815764]), layer_3 = (weight = Float32[-0.6772371 -0.19355826 … 0.092198014 -0.33821836; -0.29864177 -0.09485075 … 0.022576144 -0.17590499], bias = Float32[-1.1515998, -0.55646694])) ``` Now we will compile the gradient function using `Reactant.@compile`. ```julia function enzyme_gradient(model, ps, st, x, y) return Enzyme.gradient(Enzyme.Reverse, Const(loss_function), Const(model), ps, Const(st), Const(x), Const(y))[2] end enzyme_gradient_compiled = @compile enzyme_gradient(model, ps_ra, st_ra, x_ra, y_ra) ∂ps_enzyme = enzyme_gradient_compiled(model, ps_ra, st_ra, x_ra, y_ra) ``` ```ansi (layer_1 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[0.2601668 -0.092877164; -0.028029911 -0.013659924; … ; -0.07384163 -0.060032368; 0.042984392 0.05160542]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[0.12923914, -0.009405826, -0.026362803, -0.014524885, 0.013915381, 0.09343623, 0.08193635, 0.0077627855, 0.001044217, 0.018755728 … 0.06704106, -0.043209504, 0.104868725, 0.014353443, 0.024228811, -0.065829255, 0.010303015, 0.09878268, 0.0678494, -0.08268405])), layer_2 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.0044578887 0.00021804239 … -0.0033274866 -0.008014374; 0.13051207 -0.0046890327 … 0.038353655 0.0933025; … ; -0.041253403 0.00088798575 … 0.0007470111 -0.020347355; 0.063391745 -0.0021197703 … 0.012145702 0.068770304]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-0.006240551, 0.13950492, -0.22439212, -0.113269635, -0.02316084, 0.14702772, 0.035196126, 0.1398194, -0.23715453, 0.32662553 … -0.014224287, 0.009401781, 0.18295963, 0.13164552, 0.16955197, -0.110567965, -0.0074349036, 0.118868664, -0.026588855, 0.031815786])), layer_3 = (weight = Reactant.ConcretePJRTArray{Float32, 2, 1}(Float32[-0.6772371 -0.19355828 … 0.092198 -0.33821842; -0.29864174 -0.09485076 … 0.022576137 -0.1759051]), bias = Reactant.ConcretePJRTArray{Float32, 1, 1}(Float32[-1.1515996, -0.55646694]))) ``` Now we check the difference: ```julia fmap(Broadcast.BroadcastFunction(-), ∂ps_zyg, ∂ps_enzyme |> cpu_device()) ``` ```ansi (layer_1 = (weight = Float32[-5.9604645f-8 -1.4901161f-8; 1.8626451f-9 -2.7939677f-9; … ; -7.450581f-9 -2.2351742f-8; 2.2351742f-8 -3.7252903f-9], bias = Float32[0.0, -9.313226f-10, 3.7252903f-9, -9.313226f-9, 4.656613f-9, -5.9604645f-8, 7.450581f-9, 1.8626451f-9, 9.313226f-10, 1.8626451f-8 … -7.450581f-9, 0.0, -7.450581f-9, -4.656613f-9, -3.7252903f-9, -1.4901161f-8, -1.8626451f-9, 1.4901161f-8, 1.4901161f-8, -1.4901161f-8]), layer_2 = (weight = Float32[-4.656613f-10 -2.910383f-11 … 4.656613f-10 -9.313226f-10; -2.9802322f-8 -9.313226f-10 … 7.450581f-9 -7.450581f-9; … ; 0.0 -7.566996f-10 … 8.731149f-10 1.8626451f-9; -1.4901161f-8 2.0954758f-9 … -1.4901161f-8 -3.7252903f-8], bias = Float32[-4.656613f-10, 0.0, 0.0, 0.0, 1.8626451f-9, 1.4901161f-8, -3.7252903f-9, -1.4901161f-8, -4.4703484f-8, 2.9802322f-8 … -1.8626451f-9, -5.5879354f-9, -1.4901161f-8, 0.0, 0.0, 0.0, 8.381903f-9, 0.0, 1.8626451f-9, -2.2351742f-8]), layer_3 = (weight = Float32[0.0 1.4901161f-8 … 1.4901161f-8 5.9604645f-8; -2.9802322f-8 1.4901161f-8 … 7.450581f-9 1.0430813f-7], bias = Float32[-1.1920929f-7, 0.0])) ``` ## Using the `TrainState` API {#compile\_lux\_model\_trainstate} Now that we saw the low-level API let's see how to train the model without any of this boilerplate. Simply follow the following steps: 1. Create a device using `reactant_device`. Remember to load `Reactant.jl` before doing this. 2. Similar to other device functions move the model, parameters, states and data to the device. Note that you might want to use [`DeviceIterator`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.DeviceIterator) to move the data loader to the device with an iterator. 3. Construct a `TrainState` using [`Training.TrainState`](/api/Lux/utilities#Lux.Training.TrainState). 4. And most importantly use `AutoEnzyme`/`AutoReactant` while calling [`Training.single_train_step!`](/api/Lux/utilities#Lux.Training.single_train_step!) or [`Training.single_train_step`](/api/Lux/utilities#Lux.Training.single_train_step). ```julia model = Chain( Dense(2 => 4, gelu), Dense(4 => 4, gelu), Dense(4 => 2) ) ps, st = Lux.setup(Random.default_rng(), model) x_ra = [randn(Float32, 2, 32) for _ in 1:32] y_ra = [xᵢ .^ 2 for xᵢ in x_ra] ps_ra = ps |> xdev st_ra = st |> xdev dataloader = DeviceIterator(xdev, zip(x_ra, y_ra)) function train_model(model, ps, st, dataloader) train_state = Training.TrainState(model, ps, st, Adam(0.001f0)) for iteration in 1:1000 for (i, (xᵢ, yᵢ)) in enumerate(dataloader) _, loss, _, train_state = Training.single_train_step!( AutoEnzyme(), MSELoss(), (xᵢ, yᵢ), train_state ) if (iteration % 100 == 0 || iteration == 1) && i == 1 @printf("Iter: [%4d/%4d]\tLoss: %.8f\n", iteration, 1000, loss) end end end return train_state end train_model(model, ps_ra, st_ra, dataloader) ``` ```ansi Iter: [ 1/1000] Loss: 13.22820091 Iter: [ 100/1000] Loss: 2.58897161 Iter: [ 200/1000] Loss: 1.14364433 Iter: [ 300/1000] Loss: 0.37711889 Iter: [ 400/1000] Loss: 0.13413195 Iter: [ 500/1000] Loss: 0.05696166 Iter: [ 600/1000] Loss: 0.03033140 Iter: [ 700/1000] Loss: 0.01917107 Iter: [ 800/1000] Loss: 0.01335674 Iter: [ 900/1000] Loss: 0.01001781 Iter: [1000/1000] Loss: 0.00796735 ``` --- --- url: /dev/manual/exporting_to_jax.md --- # Exporting Lux Models to Jax (via EnzymeJAX & Reactant) {#exporting\_to\_stablehlo} In this manual, we will go over how to export Lux models to StableHLO and use [EnzymeJAX](https://github.com/EnzymeAD/Enzyme-JAX) to run integrate Lux models with JAX. We assume that users are familiar with [Reactant compilation of Lux models](/manual/compiling_lux_models#reactant-compilation). ```julia using Lux, Reactant, Random, NPZ const dev = reactant_device() ``` ```ansi (::ReactantDevice{Missing, Missing, Missing, Missing, Union{}}) (generic function with 1 method) ``` We simply define a Lux model and generate the stablehlo code using `Reactant.@code_hlo`. ```julia model = Chain( Conv((5, 5), 1 => 6, relu), MaxPool((2, 2)), Conv((5, 5), 6 => 16, relu), MaxPool((2, 2)), FlattenLayer(3), Chain( Dense(256 => 128, relu), Dense(128 => 84, relu), Dense(84 => 10) ) ) ps, st = Lux.setup(Random.default_rng(), model) |> dev ``` Generate an example input. ```julia x = randn(Random.default_rng(), Float32, 28, 28, 1, 4) |> dev ``` Now instead of compiling the model, we will use the `Reactant.Serialization.export_to_enzymejax` function to export the model. ```julia python_file_path = Reactant.Serialization.export_to_enzymejax( model, x, ps, st; function_name="run_lux_model" ) ``` ```ansi "/tmp/jl_i4ytvk/run_lux_model.py" ``` This will generate a python file that can be used to run the model using EnzymeJAX. ```julia println(read(open(python_file_path, "r"), String)) ``` ```ansi """ Auto-generated Python script for calling exported Julia/Reactant function via EnzymeJAX. This script was generated by Reactant.Serialization.export_to_enzymejax(). """ from enzyme_ad.jax import hlo_call import jax from jax.sharding import PartitionSpec as P import jax.numpy as jnp import numpy as np import os # Get the directory of this script _script_dir = os.path.dirname(os.path.abspath(__file__)) # Load the MLIR/StableHLO code with open(os.path.join(_script_dir, "run_lux_model_0.mlir"), "r") as f: _hlo_code = f.read() def load_inputs(): """Load the example inputs that were exported from Julia.""" npz_data = np.load(os.path.join(_script_dir, "run_lux_model_0_inputs.npz")) inputs = [npz_data['arr_1'], npz_data['arr_2'], npz_data['arr_3'], npz_data['arr_4'], npz_data['arr_5'], npz_data['arr_6'], npz_data['arr_7'], npz_data['arr_8'], npz_data['arr_9'], npz_data['arr_10'], npz_data['arr_11']] return tuple(inputs) @jax.jit def run_run_lux_model(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11): """ Call the exported Julia function via EnzymeJAX. Args: arg1: Array of shape (4, 1, 28, 28) and dtype float32. Path: arg.2 arg2: Array of shape (6, 1, 5, 5) and dtype float32. Path: arg.3.1.1 arg3: Array of shape (6,) and dtype float32. Path: arg.3.1.2 arg4: Array of shape (16, 6, 5, 5) and dtype float32. Path: arg.3.3.1 arg5: Array of shape (16,) and dtype float32. Path: arg.3.3.2 arg6: Array of shape (256, 128) and dtype float32. Path: arg.3.6.1.1 arg7: Array of shape (128,) and dtype float32. Path: arg.3.6.1.2 arg8: Array of shape (128, 84) and dtype float32. Path: arg.3.6.2.1 arg9: Array of shape (84,) and dtype float32. Path: arg.3.6.2.2 arg10: Array of shape (84, 10) and dtype float32. Path: arg.3.6.3.1 arg11: Array of shape (10,) and dtype float32. Path: arg.3.6.3.2 Returns: The result of calling the exported function. Note: All inputs must be in row-major (Python/NumPy) order. If you're passing arrays from Julia, make sure to transpose them first using: `permutedims(arr, reverse(1:ndims(arr)))` """ assert arg1.dtype == np.dtype('float32'), f"Expected dtype of arg1 to be float32. Got {arg1.dtype} (path: arg.2)" assert arg2.dtype == np.dtype('float32'), f"Expected dtype of arg2 to be float32. Got {arg2.dtype} (path: arg.3.1.1)" assert arg3.dtype == np.dtype('float32'), f"Expected dtype of arg3 to be float32. Got {arg3.dtype} (path: arg.3.1.2)" assert arg4.dtype == np.dtype('float32'), f"Expected dtype of arg4 to be float32. Got {arg4.dtype} (path: arg.3.3.1)" assert arg5.dtype == np.dtype('float32'), f"Expected dtype of arg5 to be float32. Got {arg5.dtype} (path: arg.3.3.2)" assert arg6.dtype == np.dtype('float32'), f"Expected dtype of arg6 to be float32. Got {arg6.dtype} (path: arg.3.6.1.1)" assert arg7.dtype == np.dtype('float32'), f"Expected dtype of arg7 to be float32. Got {arg7.dtype} (path: arg.3.6.1.2)" assert arg8.dtype == np.dtype('float32'), f"Expected dtype of arg8 to be float32. Got {arg8.dtype} (path: arg.3.6.2.1)" assert arg9.dtype == np.dtype('float32'), f"Expected dtype of arg9 to be float32. Got {arg9.dtype} (path: arg.3.6.2.2)" assert arg10.dtype == np.dtype('float32'), f"Expected dtype of arg10 to be float32. Got {arg10.dtype} (path: arg.3.6.3.1)" assert arg11.dtype == np.dtype('float32'), f"Expected dtype of arg11 to be float32. Got {arg11.dtype} (path: arg.3.6.3.2)" assert arg1.shape == (4, 1, 28, 28), f"Expected shape of arg1 to be (4, 1, 28, 28). Got {arg1.shape} (path: arg.2)" assert arg2.shape == (6, 1, 5, 5), f"Expected shape of arg2 to be (6, 1, 5, 5). Got {arg2.shape} (path: arg.3.1.1)" assert arg3.shape == (6,), f"Expected shape of arg3 to be (6,). Got {arg3.shape} (path: arg.3.1.2)" assert arg4.shape == (16, 6, 5, 5), f"Expected shape of arg4 to be (16, 6, 5, 5). Got {arg4.shape} (path: arg.3.3.1)" assert arg5.shape == (16,), f"Expected shape of arg5 to be (16,). Got {arg5.shape} (path: arg.3.3.2)" assert arg6.shape == (256, 128), f"Expected shape of arg6 to be (256, 128). Got {arg6.shape} (path: arg.3.6.1.1)" assert arg7.shape == (128,), f"Expected shape of arg7 to be (128,). Got {arg7.shape} (path: arg.3.6.1.2)" assert arg8.shape == (128, 84), f"Expected shape of arg8 to be (128, 84). Got {arg8.shape} (path: arg.3.6.2.1)" assert arg9.shape == (84,), f"Expected shape of arg9 to be (84,). Got {arg9.shape} (path: arg.3.6.2.2)" assert arg10.shape == (84, 10), f"Expected shape of arg10 to be (84, 10). Got {arg10.shape} (path: arg.3.6.3.1)" assert arg11.shape == (10,), f"Expected shape of arg11 to be (10,). Got {arg11.shape} (path: arg.3.6.3.2)" return hlo_call( arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11, source=_hlo_code, ) if __name__ == "__main__": # Load the example inputs (arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11,) = load_inputs() # Convert inputs to jax arrays arg1 = jnp.asarray(arg1) arg2 = jnp.asarray(arg2) arg3 = jnp.asarray(arg3) arg4 = jnp.asarray(arg4) arg5 = jnp.asarray(arg5) arg6 = jnp.asarray(arg6) arg7 = jnp.asarray(arg7) arg8 = jnp.asarray(arg8) arg9 = jnp.asarray(arg9) arg10 = jnp.asarray(arg10) arg11 = jnp.asarray(arg11) # Run the function print("Running run_lux_model...") result = run_run_lux_model(arg1, arg2, arg3, arg4, arg5, arg6, arg7, arg8, arg9, arg10, arg11) print("Result:", result) ``` --- --- url: /dev/manual/profiling_training_loop.md --- # Profiling Lux Training Loops {#profiling-training-loop-reactant} ::: warning Only for Reactant This tutorial is applicable iff you are using `Reactant.jl` (`AutoEnzyme` with `ReactantDevice`) for training. ::: To profile the training loop, wrap the training loop with `Reactant.with_profiler` and pass the path to the directory where the traces should be saved. Note that this will have some overhead and hence should be used only for debugging purposes. A simple example is shown below: ```julia using Reactant, Lux, Random, MLUtils, Optimisers dev = reactant_device() x_data = rand(Float32, 32, 1024) y_data = x_data .^ 2 .- 1 dl = DataLoader((x_data, y_data); batchsize=32, shuffle=true) |> dev; model = Chain(Dense(32 => 64, relu), Dense(64 => 32)) ps, st = Lux.setup(Random.default_rng(), model) |> dev; Reactant.with_profiler(joinpath(tempdir(), "lux_training_trace")) do train_state = Training.TrainState(model, ps, st, Adam(0.001)) for epoch in 1:10 for (x, y) in dl _, loss, _, train_state = Training.single_train_step!( AutoEnzyme(), MSELoss(), (x, y), train_state; return_gradients=Val(false) ) end end end ``` ```ansi WARNING: All log messages before absl::InitializeLog() is called are written to STDERR I0000 00:00:1780291933.439303 9225 profiler_session.cc:119] Profiler session initializing. I0000 00:00:1780291933.439337 9225 profiler_session.cc:134] Profiler session started. I0000 00:00:1780291965.911461 9225 profiler_session.cc:82] Profiler session collecting data. I0000 00:00:1780291965.977010 9225 save_profile.cc:150] Collecting XSpace to repository: /tmp/lux_training_trace/plugins/profile/2026_06_01_05_32_45/runnervm3jyl0.xplane.pb I0000 00:00:1780291966.033818 9225 save_profile.cc:123] Creating directory: /tmp/lux_training_trace/plugins/profile/2026_06_01_05_32_45 I0000 00:00:1780291966.084137 9225 save_profile.cc:129] Dumped gzipped tool data for trace.json.gz to /tmp/lux_training_trace/plugins/profile/2026_06_01_05_32_45/runnervm3jyl0.trace.json.gz I0000 00:00:1780291966.107417 9225 profiler_session.cc:152] Profiler session tear down. ``` Once the run is completed, you can use [`xprof`](https://github.com/openxla/xprof) to analyze the traces. An example of the output is shown below: --- --- url: /dev/manual/autodiff.md --- # Automatic Differentiation {#autodiff-lux} Lux is not an AD package, but it composes well with most of the AD packages available in the Julia ecosystem. This document lists the current level of support for various AD packages in Lux. Additionally, we provide some convenience functions for working with AD. ## Overview {#Overview} | AD Package | Mode | CPU | GPU | TPU | Nested 2nd Order AD | Support Class | |:-------------------------------------------------------------------------------------------------------------------- |:----------------- |:----- |:----- |:--- |:------------------- |:------------- | | [`Reactant.jl`](https://github.com/EnzymeAD/Reactant.jl)\[^re] + [`Enzyme.jl`](https://github.com/EnzymeAD/Enzyme.jl) | Reverse / Forward | ✔️ | ✔️ | ✔️ | ✔️ | Tier I | | [`ChainRules.jl`](https://github.com/JuliaDiff/ChainRules.jl)\[^cr] | Reverse | ✔️ | ✔️ | ❌ | ✔️ | Tier I | | [`Enzyme.jl`](https://github.com/EnzymeAD/Enzyme.jl) | Reverse / Forward | ✔️ | ❓\[^q] | ❌ | ❓\[^q] | Tier I\[^e] | | [`Zygote.jl`](https://github.com/FluxML/Zygote.jl) | Reverse | ✔️ | ✔️ | ❌ | ✔️ | Tier I | | [`ForwardDiff.jl`](https://github.com/JuliaDiff/ForwardDiff.jl) | Forward | ✔️ | ✔️ | ❌ | ✔️ | Tier I | | [`Mooncake.jl`](https://github.com/compintell/Mooncake.jl) | Reverse | ❓\[^q] | ❌ | ❌ | ❌ | Tier III | | [`ReverseDiff.jl`](https://github.com/JuliaDiff/ReverseDiff.jl) | Reverse | ✔️ | ❌ | ❌ | ❌ | Tier IV | | [`Tracker.jl`](https://github.com/FluxML/Tracker.jl) | Reverse | ✔️ | ✔️ | ❌ | ❌ | Tier IV | | [`Diffractor.jl`](https://github.com/JuliaDiff/Diffractor.jl) | Forward | ❓\[^q] | ❓\[^q] | ❌ | ❓\[^q] | Tier IV | \[^e]: Currently Enzyme outperforms other AD packages in terms of CPU performance. However, there are some edge cases where it might not work with Lux when not using Reactant. We are working on improving the compatibility. Please report any issues you encounter and try Reactant if something fails. \[^q]: This feature is supported downstream, but we don't extensively test it to ensure that it works with Lux. \[^cr]: Note that `ChainRules.jl` is not really an AD package, but we have first-class support for packages that use `rrules`. \[^re]: Note that `Reactant.jl` is not really an AD package, but a tool for compiling functions, including the use of EnzymeMLIR for AD via `Enzyme.jl`. We have first-class support for the usage of `Reactant.jl` for inference and training when using `Enzyme.jl` for differentiation. ## Recommendations {#autodiff-recommendations} * For CPU Use cases: 1. Use `Reactant.jl` + `Enzyme.jl` for the best performance as well as mutation-support. When available, this is the most reliable and fastest option. 2. Use `Zygote.jl` for the best performance without `Reactant.jl`. This is the most reliable and fastest option for CPU for the time-being. (We are working on faster Enzyme support for CPU) 3. Use `Enzyme.jl`, if there are mutations in the code and/or `Zygote.jl` fails. 4. If `Enzyme.jl` fails for some reason, (open an issue and) try `ReverseDiff.jl` ([possibly with compiled mode](https://juliadiff.org/ReverseDiff.jl/dev/api/#ReverseDiff.compile)). * For GPU Use cases: 1. Use `Reactant.jl` + `Enzyme.jl` for the best performance. This is the most reliable and fastest option for both NVIDIA and AMD GPUs. 2. Use `Zygote.jl` for the best performance on non-NVIDIA GPUs if not using Reactant. This is the most reliable and fastest non-`Reactant.jl` option for GPU for the time-being. We are working on supporting `Enzyme.jl` without `Reactant.jl` for GPU as well. * For TPU Use cases: 1. Use `Reactant.jl`. This is the only supported (and fastest) option. ## Support Class {#Support-Class} 1. **Tier I**: These packages are fully supported and have been tested extensively. Often have special rules to enhance performance. Issues for these backends take the highest priority. 2. **Tier II**: These packages are supported and extensively tested but often don't have the best performance. Issues against these backends are less critical, but we fix them when possible. (Some specific edge cases, especially with AMDGPU, are known to fail here) 3. **Tier III**: We don't know if these packages currently work with Lux. We'd love to add tests for these backends, but currently these are not our priority. 4. **Tier IV**: At some point we may or may not have supported these frameworks. However, these are not maintained at the moment. Whatever code exists for these frameworks in Lux, may be removed in the future (with a breaking release). It is not recommended to use these frameworks with Lux. ## Footnotes {#Footnotes} --- --- url: /dev/manual/nested_autodiff.md --- # Nested Automatic Differentiation {#nested\_autodiff} ::: tip Reactant Reactant.jl natively supports nested AD (with orders greater than 2nd order). For more robust nested AD, use Lux with Reactant.jl. ::: In this manual, we will explore how to use automatic differentiation (AD) inside your layers or loss functions and have Lux automatically switch the AD backend with a faster one when needed. ::: tip Tip Don't wan't Lux to do this switching for you? You can disable it by setting the `automatic_nested_ad_switching` Preference to `false`. Remember that if you are using ForwardDiff inside a Zygote call, it will drop gradients (with a warning message), so it is not recommended to use this combination. ::: Let's explore this using some questions that were posted on the [Julia Discourse forum](https://discourse.julialang.org/). ```julia using ADTypes, Lux, LinearAlgebra, Zygote, ForwardDiff, Random, StableRNGs using ComponentArrays, FiniteDiff ``` First let's set the stage using some minor changes that need to be made for this feature to work: * Switching only works if a [`StatefulLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer) is being used, with the following function calls: * For operations on the inputs: * `( ∘ )(x::AbstractArray)` * `( ∘ )(x::AbstractArray)` * `()(x::AbstractArray)` * For operations on the parameters: * `( ∘ Base.Fix1(, x))(ps)` * `(Base.Fix1(, x) ∘ )(ps)` * `(Base.Fix1(, x))(ps)` * Currently we have custom routines implemented for: * `Zygote.` * `ForwardDiff.` * [`vector_jacobian_product`](/api/Lux/autodiff#Lux.vector_jacobian_product) * [`jacobian_vector_product`](/api/Lux/autodiff#Lux.jacobian_vector_product) * [`batched_jacobian`](/api/Lux/autodiff#Lux.batched_jacobian) * Switching only happens for `ChainRules` compatible AD libraries. We plan to capture `DifferentiationInterface`, and `Enzyme.autodiff` calls in the future (PRs are welcome). ::: tip Tip [`@compact`](/api/Lux/utilities#Lux.@compact) uses [`StatefulLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer)s internally, so you can directly use these features inside a layer generated by [`@compact`](/api/Lux/utilities#Lux.@compact). ::: ## Loss Function containing Jacobian Computation {#Loss-Function-containing-Jacobian-Computation} This problem comes from `@facusapienza` on [Discourse](https://discourse.julialang.org/t/nested-and-different-ad-methods-altogether-how-to-add-ad-calculations-inside-my-loss-function-when-using-neural-differential-equations/108985). In this case, we want to add a regularization term to the neural DE based on first-order derivatives. The neural DE part is not important here and we can demonstrate this easily with a standard neural network. ```julia function loss_function1(model, x, ps, st, y) # Make it a stateful layer smodel = StatefulLuxLayer(model, ps, st) ŷ = smodel(x) loss_emp = sum(abs2, ŷ .- y) # You can use `Zygote.jacobian` as well but ForwardDiff tends to be more efficient here J = ForwardDiff.jacobian(smodel, x) loss_reg = abs2(norm(J .* 0.01f0)) return loss_emp + loss_reg end # Using Batchnorm to show that it is possible model = Chain(Dense(2 => 4, tanh), BatchNorm(4), Dense(4 => 2)) ps, st = Lux.setup(StableRNG(0), model) x = randn(StableRNG(0), Float32, 2, 10) y = randn(StableRNG(11), Float32, 2, 10) loss_function1(model, x, ps, Lux.testmode(st), y) ``` ```ansi 11.380776f0 ``` So our loss function works, let's take the gradient (forward diff doesn't nest nicely here): ```julia _, ∂x, ∂ps, _, _ = Zygote.gradient(loss_function1, model, x, ps, st, y) ``` ```ansi (nothing, Float32[-1.6702257 0.9043228 … 0.16094846 -4.992662; -8.010404 0.8541596 … 3.3928175 -7.1936812], (layer_1 = (weight = Float32[-4.3707023 -4.9076533; 22.199387 1.867202; 0.47872233 -0.9734574; -0.36428708 0.31861955], bias = Float32[-1.0168695, -0.16566901, 1.0829282, 1.4810884]), layer_2 = (scale = Float32[4.2774315, 3.1984668, 6.840588, 3.7018592], bias = Float32[-2.6477456, 4.9094505, -4.987689, -0.7292344]), layer_3 = (weight = Float32[11.395306 1.9206433 9.744489 -7.6726513; 2.5979974 7.106069 -7.869632 -1.787159], bias = Float32[0.041031003, 7.928609])), nothing, Float32[0.48193252 1.4007905 … -0.19124654 -1.7181164; 1.7811481 0.6913705 … -1.5627227 1.4397957]) ``` Now let's verify the gradients using finite differences: ```julia ∂x_fd = FiniteDiff.finite_difference_gradient(x -> loss_function1(model, x, ps, st, y), x) ∂ps_fd = FiniteDiff.finite_difference_gradient(ps -> loss_function1(model, x, ps, st, y), ComponentArray(ps)) println("∞-norm(∂x - ∂x_fd): ", norm(∂x .- ∂x_fd, Inf)) println("∞-norm(∂ps - ∂ps_fd): ", norm(ComponentArray(∂ps) .- ∂ps_fd, Inf)) ``` ```ansi ┌ Warning: `training` is set to `Val{true}()` but is not being used within an autodiff call (gradient, jacobian, etc...). This will be slow. If you are using a `Lux.jl` model, set it to inference (test) mode using `LuxCore.testmode`. Reliance on this behavior is discouraged, and is not guaranteed by Semantic Versioning, and might be removed without a deprecation cycle. It is recommended to fix this issue in your code. └ @ LuxLib.Utils ~/work/Lux.jl/Lux.jl/lib/LuxLib/src/utils.jl:346 ┌ Warning: `training` is set to `Val{true}()` but is not being used within an autodiff call (gradient, jacobian, etc...). This will be slow. If you are using a `Lux.jl` model, set it to inference (test) mode using `LuxCore.testmode`. Reliance on this behavior is discouraged, and is not guaranteed by Semantic Versioning, and might be removed without a deprecation cycle. It is recommended to fix this issue in your code. └ @ LuxLib.Utils ~/work/Lux.jl/Lux.jl/lib/LuxLib/src/utils.jl:346 ∞-norm(∂x - ∂x_fd): 0.00046014786 ∞-norm(∂ps - ∂ps_fd): 0.00068473816 ``` That's pretty good, of course you will have some error from the finite differences calculation. ### Using Batched Jacobian for Multiple Inputs {#Using-Batched-Jacobian-for-Multiple-Inputs} Notice that in this example the Jacobian `J` consists on the full matrix of derivatives of `smodel` with respect the different inputs in `x`. In many cases, we are interested in computing the Jacobian with respect to each input individually, avoiding the unnecessary calculation of zero entries of the Jacobian. This can be achieved with [`batched_jacobian`](/api/Lux/autodiff#Lux.batched_jacobian) to parse the calculation of the Jacobian per each single input. Using the same example from the previous section: ```julia model = Chain(Dense(2 => 4, tanh), Dense(4 => 2)) ps, st = Lux.setup(StableRNG(0), model) x = randn(StableRNG(0), Float32, 2, 10) y = randn(StableRNG(11), Float32, 2, 10) function loss_function_batched(model, x, ps, st, y) # Make it a stateful layer smodel = StatefulLuxLayer(model, ps, st) ŷ = smodel(x) loss_emp = sum(abs2, ŷ .- y) # You can use `AutoZygote()` as well but `AutoForwardDiff()` tends to be more efficient here J = batched_jacobian(smodel, AutoForwardDiff(), x) loss_reg = abs2(norm(J .* 0.01f0)) return loss_emp + loss_reg end loss_function_batched(model, x, ps, st, y) ``` ```ansi 11.380777f0 ``` Notice that in this last example we removed `BatchNorm()` from the neural network. This is done so outputs corresponding to different inputs don't have an algebraic dependency due to the batch normalization happening in the neural network. We can now verify again the value of the Jacobian: ```julia ∂x_fd = FiniteDiff.finite_difference_gradient(x -> loss_function_batched(model, x, ps, st, y), x) ∂ps_fd = FiniteDiff.finite_difference_gradient(ps -> loss_function_batched(model, x, ps, st, y), ComponentArray(ps)) _, ∂x_b, ∂ps_b, _, _ = Zygote.gradient(loss_function_batched, model, x, ps, st, y) println("∞-norm(∂x_b - ∂x_fd): ", norm(∂x_b .- ∂x_fd, Inf)) println("∞-norm(∂ps_b - ∂ps_fd): ", norm(ComponentArray(∂ps_b) .- ∂ps_fd, Inf)) ``` ```ansi ∞-norm(∂x_b - ∂x_fd): 0.00020849705 ∞-norm(∂ps_b - ∂ps_fd): 0.00025326014 ``` In this example, it is important to remark that now `batched_jacobian` returns a 3D array with the Jacobian calculation for each independent input value in `x`. ## Loss Function contains Gradient Computation {#Loss-Function-contains-Gradient-Computation} Ok here I am going to cheat a bit. This comes from a discussion on nested AD for PINNs on [Discourse](https://discourse.julialang.org/t/is-it-possible-to-do-nested-ad-elegantly-in-julia-pinns/98888/21). As the consensus there, we shouldn't use nested AD for 3rd or higher order differentiation. Note that in the example there, the user uses `ForwardDiff.derivative` but we will use `ForwardDiff.gradient` instead, as we typically deal with array inputs and outputs. ```julia function loss_function2(model, t, ps, st) smodel = StatefulLuxLayer(model, ps, st) ŷ = only(Zygote.gradient(Base.Fix1(sum, abs2) ∘ smodel, t)) # Zygote returns a tuple return sum(abs2, ŷ .- cos.(t)) end model = Chain(Dense(1 => 12,tanh), Dense(12 => 12,tanh), Dense(12 => 12,tanh), Dense(12 => 1)) ps, st = Lux.setup(StableRNG(0), model) t = rand(StableRNG(0), Float32, 1, 16) ``` ```ansi 1×16 Matrix{Float32}: 0.420698 0.488105 0.267644 0.784768 … 0.305844 0.131726 0.859405 ``` Now the moment of truth: ```julia _, ∂t, ∂ps, _ = Zygote.gradient(loss_function2, model, t, ps, st) ``` ```ansi (nothing, Float32[-0.55306894 0.15707001 … -8.553633 0.07513529], (layer_1 = (weight = Float32[-1.3108878; -2.4101036; … ; 0.43676856; 1.9626999;;], bias = Float32[-1.77037, 1.7834253, -7.107933, -3.4437156, 3.2615936, -1.9511774, 11.527171, -1.8003641, 6.7513776, -4.77004, -3.183308, 6.587816]), layer_2 = (weight = Float32[-0.23921251 -0.20668766 … -0.6383875 -2.23242; -1.6666821 1.042543 … -1.640935 -3.4007297; … ; -0.36023283 -0.086430185 … -0.7054552 -2.1921258; 3.1173706 -1.9727281 … 3.0402095 6.11373], bias = Float32[0.3729237, -2.9340098, 3.6637254, -0.72471243, -0.7925039, -1.1245009, -0.89859, -0.03284671, -2.729647, -8.446214, 0.06208062, 5.5367613]), layer_3 = (weight = Float32[-0.7262076 1.0381727 … -1.5016018 -1.679885; 2.2896144 0.43350345 … -1.6663245 -1.8067697; … ; -2.1851237 -0.64241976 … 1.9577401 2.148901; 0.3654292 -0.09699093 … 0.025357665 0.028738922], bias = Float32[1.1350523, -2.1769388, 4.1149755, 3.2842, 0.35638645, 3.7911117, -0.007984848, -2.0338566, -1.1642132, -2.9500434, 2.0285957, -0.41238895]), layer_4 = (weight = Float32[15.794909 -20.65178 … -7.798029 -9.910249], bias = Float32[11.461399])), nothing) ``` Boom that worked! Let's verify the gradient using forward diff: ```julia ∂t_fd = ForwardDiff.gradient(t -> loss_function2(model, t, ps, st), t) ∂ps_fd = ForwardDiff.gradient(ps -> loss_function2(model, t, ps, st), ComponentArray(ps)) println("∞-norm(∂t - ∂t_fd): ", norm(∂t .- ∂t_fd, Inf)) println("∞-norm(∂ps - ∂ps_fd): ", norm(ComponentArray(∂ps) .- ∂ps_fd, Inf)) ``` ```ansi ∞-norm(∂t - ∂t_fd): 5.722046e-6 ∞-norm(∂ps - ∂ps_fd): 2.861023e-6 ``` ## Loss Function computing the Jacobian of the Parameters {#Loss-Function-computing-the-Jacobian-of-the-Parameters} The above example shows how to compute the gradient/jacobian wrt the inputs in the loss function. However, what if we want to compute the jacobian wrt the parameters? This problem has been taken from [Issue 610](https://github.com/LuxDL/Lux.jl/issues/610). We resolve these setups by using the `Base.Fix1` wrapper around the stateful layer and fixing the input to the stateful layer. ```julia function loss_function3(model, x, ps, st) smodel = StatefulLuxLayer(model, ps, st) J = only(Zygote.jacobian(Base.Fix1(smodel, x), ps)) # Zygote returns a tuple return sum(abs2, J) end model = Chain(Dense(1 => 12,tanh), Dense(12 => 12,tanh), Dense(12 => 12,tanh), Dense(12 => 1)) ps, st = Lux.setup(StableRNG(0), model) ps = ComponentArray(ps) # needs to be an AbstractArray for most jacobian functions x = rand(StableRNG(0), Float32, 1, 16) ``` ```ansi 1×16 Matrix{Float32}: 0.420698 0.488105 0.267644 0.784768 … 0.305844 0.131726 0.859405 ``` We can as usual compute the gradient/jacobian of the loss function: ```julia _, ∂x, ∂ps, _ = Zygote.gradient(loss_function3, model, x, ps, st) ``` ```ansi (nothing, Float32[6.8464594 6.2111297 … 1.9693907 -1.959184], (layer_1 = (weight = Float32[-3.6867144; -1.68539; … ; 2.9501405; -6.637219;;], bias = Float32[-6.488623, -7.0661273, 1.3344336, 2.6049256, 0.7290931, -15.730944, -5.431456, 7.4604826, -1.186449, 15.522138, 0.44571495, -15.376386]), layer_2 = (weight = Float32[0.3980046 -4.3071294 … -1.0914631 -4.759412; 0.8852217 -2.2523668 … 0.3977316 0.1306752; … ; -2.2192001 0.88214755 … -0.55989707 1.3939897; -3.154516 4.594261 … -1.7649312 -0.38241944], bias = Float32[7.5247808, 4.2529244, -17.252981, 3.260692, -7.4066525, 1.1126353, 2.8471048, 6.7544622, -9.815336, 0.18652153, -4.5365167, -10.04811]), layer_3 = (weight = Float32[1.0462949 4.899997 … 1.1557573 -2.284967; -2.3719285 8.687264 … -3.1904757 -8.841231; … ; -10.298787 -2.9139607 … -9.754746 -4.0381317; 1.2221471 -0.46878588 … 1.0469304 0.9091032], bias = Float32[2.8379912, 8.345026, 2.9214194, -2.2415926, -11.139433, -3.834073, -2.845412, -7.9164896, 4.222528, -1.2864517, 6.9338737, -1.4144737]), layer_4 = (weight = Float32[-59.44397 -12.688665 … 99.77208 -3.339081], bias = Float32[0.0])), nothing) ``` Now let's verify the gradient using forward diff: ```julia ∂x_fd = ForwardDiff.gradient(x -> loss_function3(model, x, ps, st), x) ∂ps_fd = ForwardDiff.gradient(ps -> loss_function3(model, x, ps, st), ComponentArray(ps)) println("∞-norm(∂x - ∂x_fd): ", norm(∂x .- ∂x_fd, Inf)) println("∞-norm(∂ps - ∂ps_fd): ", norm(ComponentArray(∂ps) .- ∂ps_fd, Inf)) ``` ```ansi ∞-norm(∂x - ∂x_fd): 1.9073486e-6 ∞-norm(∂ps - ∂ps_fd): 3.0517578e-5 ``` ## Hutchinson Trace Estimation {#Hutchinson-Trace-Estimation} Hutchinson Trace Estimation often shows up in machine learning literature to provide a fast estimate of the trace of a Jacobian Matrix. This is based off of [Hutchinson 1990](https://www.nowozin.net/sebastian/blog/thoughts-on-trace-estimation-in-deep-learning.html) which computes the estimated trace of a matrix $A \in \mathbb{R}^{D \times D}$ using random vectors $v \in \mathbb{R}^{D}$ s.t. $\mathbb{E}\left\[v v^T\right] = I$. $$\text{Tr}(A) = \mathbb{E}\left\[v^T A v\right] = \frac{1}{V} \sum\_{i = 1}^V v\_i^T A v\_i$$ We can use this to compute the trace of a Jacobian Matrix $J \in \mathbb{R}^{D \times D}$ using the following algorithm: $$\text{Tr}(J) = \frac{1}{V} \sum\_{i = 1}^V v\_i^T J v\_i$$ Note that we can compute this using two methods: 1. Compute $v\_i^T J$ using a Vector-Jacobian product and then do a matrix-vector product to get the trace. 2. Compute $J v\_i$ using a Jacobian-Vector product and then do a matrix-vector product to get the trace. For simplicity, we will use a single sample of $v\_i$ to compute the trace. Additionally, we will fix the sample to ensure that our tests against the finite difference implementation are not affected by the randomness in the sample. ### Computing using the Vector-Jacobian Product {#Computing-using-the-Vector-Jacobian-Product} ```julia function hutchinson_trace_vjp(model, x, ps, st, v) smodel = StatefulLuxLayer(model, ps, st) vjp = vector_jacobian_product(smodel, AutoZygote(), x, v) return sum(batched_matmul(reshape(vjp, 1, :, size(vjp, ndims(vjp))), reshape(v, :, 1, size(v, ndims(v))))) end ``` ```ansi hutchinson_trace_vjp (generic function with 1 method) ``` This vjp version will be the fastest and most scalable and hence is the recommended way for computing hutchinson trace. ### Computing using the Jacobian-Vector Product {#Computing-using-the-Jacobian-Vector-Product} ```julia function hutchinson_trace_jvp(model, x, ps, st, v) smodel = StatefulLuxLayer(model, ps, st) jvp = jacobian_vector_product(smodel, AutoForwardDiff(), x, v) return sum(batched_matmul(reshape(v, 1, :, size(v, ndims(v))), reshape(jvp, :, 1, size(jvp, ndims(jvp))))) end ``` ```ansi hutchinson_trace_jvp (generic function with 1 method) ``` ### Computing using the Full Jacobian {#Computing-using-the-Full-Jacobian} This is definitely not recommended, but we are showing it for completeness. ```julia function hutchinson_trace_full_jacobian(model, x, ps, st, v) smodel = StatefulLuxLayer(model, ps, st) J = ForwardDiff.jacobian(smodel, x) return vec(v)' * J * vec(v) end ``` ```ansi hutchinson_trace_full_jacobian (generic function with 1 method) ``` Now let's compute the trace and compare the results: ```julia model = Chain(Dense(4 => 12,tanh), Dense(12 => 12,tanh), Dense(12 => 12,tanh), Dense(12 => 4)) ps, st = Lux.setup(StableRNG(0), model) x = rand(StableRNG(0), Float32, 4, 12) v = (rand(StableRNG(12), Float32, 4, 12) .> 0.5f0) * 2.0f0 .- 1.0f0 # rademacher sample ``` ```julia tr_vjp = hutchinson_trace_vjp(model, x, ps, st, v) tr_jvp = hutchinson_trace_jvp(model, x, ps, st, v) tr_full_jacobian = hutchinson_trace_full_jacobian(model, x, ps, st, v) println("Tr(J) using vjp: ", tr_vjp) println("Tr(J) using jvp: ", tr_jvp) println("Tr(J) using full jacobian: ", tr_full_jacobian) ``` ```ansi Tr(J) using vjp: 4.9127817 Tr(J) using jvp: 4.9127817 Tr(J) using full jacobian: 4.912781 ``` Now that we have verified that the results are the same, let's try to differentiate the trace estimate. This often shows up as a regularization term in neural networks. ```julia _, ∂x_vjp, ∂ps_vjp, _, _ = Zygote.gradient(hutchinson_trace_vjp, model, x, ps, st, v) _, ∂x_jvp, ∂ps_jvp, _, _ = Zygote.gradient(hutchinson_trace_jvp, model, x, ps, st, v) _, ∂x_full_jacobian, ∂ps_full_jacobian, _, _ = Zygote.gradient(hutchinson_trace_full_jacobian, model, x, ps, st, v) ``` For sanity check, let's verify that the gradients are the same: ```julia println("∞-norm(∂x using vjp): ", norm(∂x_vjp .- ∂x_jvp, Inf)) println("∞-norm(∂ps using vjp): ", norm(ComponentArray(∂ps_vjp) .- ComponentArray(∂ps_jvp), Inf)) println("∞-norm(∂x using full jacobian): ", norm(∂x_full_jacobian .- ∂x_vjp, Inf)) println("∞-norm(∂ps using full jacobian): ", norm(ComponentArray(∂ps_full_jacobian) .- ComponentArray(∂ps_vjp), Inf)) ``` ```ansi ∞-norm(∂x using vjp): 0.0 ∞-norm(∂ps using vjp): 0.0 ∞-norm(∂x using full jacobian): 7.1525574e-7 ∞-norm(∂ps using full jacobian): 1.4305115e-6 ``` --- --- url: /dev/manual/debugging.md --- # Debugging Lux Models {#debug-lux-layers} Debugging DNNs can be very painful. Especially with the gigantic stacktraces for Lux, it is even harder to pin-point to which particular layer errored out. This page describes some useful tools that ship with Lux, that can help you debug your models. ::: tip TL;DR Simply wrap your model with `Lux.Experimental.@debug_mode`!! ::: ::: warning Don't Forget Remember to use the non Debug mode model after you finish debugging. Debug mode models are way slower. ::: Let us construct a model which has an obviously incorrect dimension. In this example, you will see how easy it is to pin-point the problematic layer. ## Incorrect Model Specification: Dimension Mismatch Problems {#Incorrect-Model-Specification:-Dimension-Mismatch-Problems} ```julia using Lux, Random model = Chain(Dense(1 => 16, relu), Chain(Dense(16 => 3), Dense(1 => 1)), BatchNorm(1)) model_debug = Lux.Experimental.@debug_mode model ``` ```ansi Chain( layer_1 = DebugLayer( layer = Dense(1 => 16, relu), # 32 parameters ), layer_2 = Chain( layer_1 = DebugLayer( layer = Dense(16 => 3), # 51 parameters ), layer_2 = DebugLayer( layer = Dense(1 => 1), # 2 parameters ), ), layer_3 = DebugLayer( layer = BatchNorm(1, affine=true, track_stats=true), # 2 parameters, plus 3 non-trainable ), )  # Total: 87 parameters,  # plus 3 states. ``` Note that we can use the parameters and states for `model` itself in `model_debug`, no need to make any changes. If you ran the original model this is the kind of error you would see: ```julia rng = Xoshiro(0) ps, st = Lux.setup(rng, model) x = randn(rng, Float32, 1, 2) try model(x, ps, st) catch e println(e) end ``` ```ansi DimensionMismatch("A has shape (1, 1) but B has shape (3, 2)") ``` Of course, this error will come with a detailed stacktrace, but it is still not very useful. Now let's try using the debug mode model: ```julia try model_debug(x, ps, st) catch e println(e) end ``` ```ansi [ Info: Input Type: Matrix{Float32} | Input Structure: (1, 2). [ Info: Running Layer: Dense(1 => 16, relu) at location KeyPath(:model, :layers, :layer_1)! [ Info: Output Type: Matrix{Float32} | Output Structure: (16, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (16, 2). [ Info: Running Layer: Dense(16 => 3) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_1)! [ Info: Output Type: Matrix{Float32} | Output Structure: (3, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (3, 2). [ Info: Running Layer: Dense(1 => 1) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2)! ┌ Error: Layer Dense(1 => 1) failed!! This layer is present at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2). └ @ Lux.Experimental ~/work/Lux.jl/Lux.jl/src/contrib/debug.jl:116 DimensionMismatch("A has shape (1, 1) but B has shape (3, 2)") ``` See now we know that `model.layers.layer_2.layers.layer_2` is the problematic layer. Let us fix that layer and see what happens: ```julia model = Chain(Dense(1 => 16, relu), Chain( Dense(16 => 3), # [!code --] Dense(16 => 1), # [!code ++] Dense(1 => 1)), BatchNorm(1)) ``` ```julia model_fixed = Chain(Dense(1 => 16, relu), Chain(Dense(16 => 1), Dense(1 => 1)), BatchNorm(1)) ps, st = Lux.setup(rng, model_fixed) model_fixed(x, ps, st) ``` ```ansi (Float32[-0.99998605 0.999986], (layer_1 = NamedTuple(), layer_2 = (layer_1 = NamedTuple(), layer_2 = NamedTuple()), layer_3 = (running_mean = Float32[0.07133968], running_var = Float32[0.971899], training = Val{true}()))) ``` Voila!! We have tracked down and fixed the problem. ## Tracking down NaNs {#Tracking-down-NaNs} Have you encountered those pesky little NaNs in your training? They are very hard to track down. We will create an artificially simulate NaNs in our model and see how we can track the offending layer. We can set `nan_check` to `:forward`, `:backward` or `:both` to check for NaNs in the debug model. (or even disable it by setting it to `:none`) ```julia model = Chain(Dense(1 => 16, relu), Chain(Dense(16 => 1), Dense(1 => 1)), BatchNorm(1)) ps, st = Lux.setup(rng, model) model_debug = Lux.Experimental.@debug_mode model nan_check=:both ``` ```ansi Chain( layer_1 = DebugLayer( layer = Dense(1 => 16, relu), # 32 parameters ), layer_2 = Chain( layer_1 = DebugLayer( layer = Dense(16 => 1), # 17 parameters ), layer_2 = DebugLayer( layer = Dense(1 => 1), # 2 parameters ), ), layer_3 = DebugLayer( layer = BatchNorm(1, affine=true, track_stats=true), # 2 parameters, plus 3 non-trainable ), )  # Total: 53 parameters,  # plus 3 states. ``` Let us set a value in the parameter to `NaN`: ```julia ps.layer_2.layer_2.weight[1, 1] = NaN ``` ```ansi NaN ``` Now let us run the model ```julia model(x, ps, st) ``` ```ansi (Float32[NaN NaN], (layer_1 = NamedTuple(), layer_2 = (layer_1 = NamedTuple(), layer_2 = NamedTuple()), layer_3 = (running_mean = Float32[NaN], running_var = Float32[NaN], training = Val{true}()))) ``` Ah as expected our output is `NaN`. But is is not very clear how to track where the first `NaN` occurred. Let's run the debug model and check: ```julia try model_debug(x, ps, st) catch e println(e) end ``` ```ansi [ Info: Input Type: Matrix{Float32} | Input Structure: (1, 2). [ Info: Running Layer: Dense(1 => 16, relu) at location KeyPath(:model, :layers, :layer_1)! [ Info: Output Type: Matrix{Float32} | Output Structure: (16, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (16, 2). [ Info: Running Layer: Dense(16 => 1) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_1)! [ Info: Output Type: Matrix{Float32} | Output Structure: (1, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (1, 2). [ Info: Running Layer: Dense(1 => 1) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2)! DomainError(Float32[NaN;;], "NaNs detected in parameters (@ KeyPath(:weight,)) of layer Dense(1 => 1) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2).") ``` And we have figured it out! The first `NaN` occurred in the parameters of `model.layers.layer_2.layers.layer_2`! But what if NaN occurs in the reverse pass! Let us define a custom layer and introduce a fake NaN in the backward pass. ```julia using ChainRulesCore, Zygote const CRC = ChainRulesCore offending_layer(x) = 2 .* x ``` ```ansi offending_layer (generic function with 1 method) ``` ```julia model = Chain(Dense(1 => 16, relu), Chain(Dense(16 => 1), offending_layer), BatchNorm(1)) ps, st = Lux.setup(rng, model) model(x, ps, st) ``` ```ansi (Float32[0.9999881 -0.9999881], (layer_1 = NamedTuple(), layer_2 = (layer_1 = NamedTuple(), layer_2 = NamedTuple()), layer_3 = (running_mean = Float32[0.0026271285], running_var = Float32[0.98396176], training = Val{true}()))) ``` Let us define a custom backward pass to introduce some NaNs: ```julia function CRC.rrule(::typeof(offending_layer), x) y = offending_layer(x) function ∇offending_layer(Δ) problematicΔ = CRC.@thunk begin Δ = CRC.unthunk(Δ) Δ[1] = NaN return Δ end return NoTangent(), problematicΔ end return y, ∇offending_layer end ``` Let us compute the gradient of the layer now: ```julia Zygote.gradient(ps -> sum(first(model(x, ps, st))), ps) ``` ```ansi ((layer_1 = (weight = Float32[0.0; 0.0; … ; 0.0; 0.0;;], bias = Float32[0.0, 0.0, 0.0, 0.0, NaN, 0.0, NaN, NaN, 0.0, 0.0, 0.0, NaN, NaN, NaN, 0.0, 0.0]), layer_2 = (layer_1 = (weight = Float32[NaN NaN … NaN NaN], bias = Float32[NaN]), layer_2 = nothing), layer_3 = (scale = Float32[0.0], bias = Float32[2.0])),) ``` Oh no!! A `NaN` is present in the gradient of `ps`. Let us run the debug model and see where the `NaN` occurred: ```julia model_debug = Lux.Experimental.@debug_mode model nan_check=:both try Zygote.gradient(ps -> sum(first(model_debug(x, ps, st))), ps) catch e println(e) end ``` ```ansi [ Info: Input Type: Matrix{Float32} | Input Structure: (1, 2). [ Info: Running Layer: Dense(1 => 16, relu) at location KeyPath(:model, :layers, :layer_1)! [ Info: Output Type: Matrix{Float32} | Output Structure: (16, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (16, 2). [ Info: Running Layer: Dense(16 => 1) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_1)! [ Info: Output Type: Matrix{Float32} | Output Structure: (1, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (1, 2). [ Info: Running Layer: WrappedFunction(offending_layer) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2)! [ Info: Output Type: Matrix{Float32} | Output Structure: (1, 2). [ Info: Input Type: Matrix{Float32} | Input Structure: (1, 2). [ Info: Running Layer: BatchNorm(1, affine=true, track_stats=true) at location KeyPath(:model, :layers, :layer_3)! [ Info: Output Type: Matrix{Float32} | Output Structure: (1, 2). DomainError(Float32[NaN 0.0], "NaNs detected in pullback output (x) of layer WrappedFunction(offending_layer) at location KeyPath(:model, :layers, :layer_2, :layers, :layer_2).") ``` And there you go our debug layer prints that the problem is in `WrappedFunction(offending_layer) at location model.layers.layer_2.layers.layer_2`! Once we fix the pullback of the layer, we will fix the NaNs. ## Conclusion {#Conclusion} In this manual section, we have discussed tracking down errors in Lux models. We have covered tracking incorrect model specifications and NaNs in forward and backward passes. However, remember that this is an **Experimental** feature, and there might be edge cases that don't work correctly. If you find any such cases, please open an issue on GitHub! --- --- url: /dev/manual/performance_pitfalls.md --- # Performance Pitfalls & How to Catch Them {#Performance-Pitfalls-and-How-to-Catch-Them} Go through the following documentations for general performance tips: 1. [Official Julia Performance Tips](https://docs.julialang.org/en/v1/manual/performance-tips/). 2. [Recommendations for selecting AD packages](/manual/autodiff#autodiff-recommendations). ::: tip Using Reactant? If you are using Lux with Reactant, general concerns with Julia performance like type-instabilities don't apply. Reactant automatically optimizes away the type instabilities. ::: ## Spurious Type-Promotion {#Spurious-Type-Promotion} Lux by-default uses Julia semantics for type-promotions, while this means that we do the "correct" numerical thing, this can often come as a surprise to users coming from a more deep learning background. For example, consider the following code: ```julia using Lux, Random rng = Xoshiro(0) model = Dense(2 => 2, gelu) ps, st = Lux.setup(rng, model) Lux.recursive_eltype((ps, st)) ``` ```ansi Float32 ``` As we can see that `ps` and `st` are structures with the highest precision being `Float32`. Now let's run the model using some random data: ```julia x = rand(rng, 2, 4) eltype(first(model(x, ps, st))) ``` ```ansi Float64 ``` Oops our output became `Float64`. This will be bad on CPUs but an absolute performance disaster on GPUs. The reason this happened is that our input `x` was `Float64`. Instead, we should have used `Float32` input: ```julia x = rand(rng, Float32, 2, 4) eltype(first(model(x, ps, st))) ``` ```ansi Float32 ``` This was easy to fix for a small model. But certain layers might incorrectly promote objects to a higher precision. This will cause a regression in performance. There are 2 recommendations to fix this or track them down: 1. Use [`Lux.Experimental.@debug_mode`](/manual/debugging#debug-lux-layers) to see which layer is causing the type-promotion. 2. Alternatively to control the global behavior of eltypes in Lux and allow it to auto-correct the precision use [`match_eltype`](/api/Lux/utilities#Lux.match_eltype) and the [`eltype_mismatch_handling`](/manual/preferences#automatic-eltypes-preference) preference. ## Scalar Indexing on GPU Arrays {#Scalar-Indexing-on-GPU-Arrays} When running code on GPUs, it is recommended to [disallow scalar indexing](https://cuda.juliagpu.org/stable/usage/workflow/#UsageWorkflowScalar). Note that this is disabled by default except in REPL. You can disable it even in REPL mode using: ```julia using GPUArraysCore GPUArraysCore.allowscalar(false) ``` ## Data Loading and Device Transfer {#Data-Loading-and-Device-Transfer} A common pattern for loading data and transferring data to GPUs looks like this: ```julia dataloader = DataLoader(dataset; parallel=true, batchsize=12) # from MLUtils.jl gdev = gpu_device() for (X, y) in dataloader X = X |> gdev y = y |> gdev # ... # do some computation # ... end ``` This is typically fast enough, but the data transfer to the device is happening in main process, not exploiting the parallelism in the dataloader. Instead, we can do this: ```julia dataloader = DataLoader(dataset; parallel=true, batchsize=12) # from MLUtils.jl gdev = gpu_device() for (X, y) in gdev(dataloader) # ... # do some computation # ... end ``` Here, `X` and `y` are on the gpu device `gdev` and the data transfer happens in the worker processes. Additionally, it behaves similar to `CuIterator` from CUDA.jl and eagerly frees the data after every iteration (this is device agnostic and works on all supported GPU backends). ## Type Instabilities {#Type-Instabilities} ::: tip Using Reactant? If you are using Lux with Reactant, type-instabilities won't affect performance at all. You can safely ignore this section. ::: `Lux.jl` is integrated with `DispatchDoctor.jl` to catch type instabilities. You can easily enable it by setting the `instability_check` preference. This will help you catch type instabilities in your code. For more information on how to set preferences, check out [`Lux.set_dispatch_doctor_preferences!`](/api/Lux/utilities#Lux.set_dispatch_doctor_preferences!). ## Faster Primitives {#Faster-Primitives} ::: tip Using Reactant? If you are using Lux with Reactant, we will automatically use optimized versions of the functions at the compiler level. You can safely ignore this section. ::: Prefer to use deep learning primitives and their fused variants from `LuxLib.jl` instead of `NNlib.jl`. Some of the alternatives are: 1. Replace `NNlib.batched_mul` with [`LuxLib.batched_matmul`](/api/NN_Primitives/LuxLib#LuxLib.API.batched_matmul). 2. Replace `NNlib.conv` with bias and activation with [`LuxLib.fused_conv_bias_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.fused_conv_bias_activation). 3. Replace `σ.(w * x .+ b)` with [`LuxLib.fused_dense_bias_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.fused_dense_bias_activation). 4. Replace uses of `σ.(x)` with [`LuxLib.fast_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.fast_activation) or [`LuxLib.fast_activation!!`](/api/NN_Primitives/LuxLib#LuxLib.API.fast_activation!!) (the latter one is often faster). 5. Replace uses of `σ.(x .+ b)` with [`LuxLib.bias_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.bias_activation) or [`LuxLib.bias_activation!!`](/api/NN_Primitives/LuxLib#LuxLib.API.bias_activation!!) (the latter one is often faster). ## Optional Dependencies for Performance {#Optional-Dependencies-for-Performance} ::: tip Using Reactant? These dependencies are not needed for Reactant. You can safely ignore this section. ::: For faster performance on CPUs load the following packages: 1. `LoopVectorization.jl` 2. `Octavian.jl` If these are available, we automatically use optimized versions of the layers. Though there are cases where this might be an issue (see [#980](https://github.com/LuxDL/Lux.jl/issues/980) and [disabling loop vectorization](/manual/preferences#disable_loop_vectorization)). --- --- url: /dev/manual/migrate_from_flux.md --- # Migrating from Flux to Lux {#migrate-from-flux} For the core library layers like [`Dense`](/api/Lux/layers#Lux.Dense), [`Conv`](/api/Lux/layers#Lux.Conv), etc. we have intentionally kept the API very similar to Flux. In most cases, replacing `using Flux` with `using Lux` should be enough to get you started. We cover the additional changes that you will have to make in the following example. :::code-group ```julia{1,7,9,11} [Lux] using Lux, Random, NNlib, Zygote model = Chain(Dense(2 => 4), BatchNorm(4, relu), Dense(4 => 2)) rng = Random.default_rng() x = randn(rng, Float32, 2, 4) ps, st = Lux.setup(rng, model) model(x, ps, st) gradient(ps -> sum(first(model(x, ps, st))), ps) ``` ```julia [Flux] using Flux, Random, NNlib, Zygote model = Chain(Dense(2 => 4), BatchNorm(4, relu), Dense(4 => 2)) rng = Random.default_rng() x = randn(rng, Float32, 2, 4) model(x) gradient(model -> sum(model(x)), model) ``` ::: ## Implementing Custom Layers {#Implementing-Custom-Layers} Flux and Lux operate under extremely different design philosophies regarding how layers should be implemented. A summary of the differences would be: * Flux stores everything in a single struct and relies on `Functors.@functor` and `Flux.trainable` to distinguish between trainable and non-trainable parameters. * Lux relies on the user to define `Lux.initialparameters` and `Lux.initialstates` to distinguish between trainable parameters (called "parameters") and non-trainable parameters (called "states"). Additionally, Lux layers define the model architecture, hence device transfer utilities like [`gpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_device), [`cpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.cpu_device), etc. cannot be applied on Lux layers, instead they need to be applied on the parameters and states. Let's work through a concrete example to demonstrate this. We will implement a very simple layer that computes $A \times B \times x$ where $A$ is not trainable and $B$ is trainable. :::code-group ```julia [Lux] using Lux, Random, NNlib, Zygote struct LuxLinear <: Lux.AbstractLuxLayer init_A init_B end function LuxLinear(A::AbstractArray, B::AbstractArray) # Storing Arrays or any mutable structure inside a Lux Layer is not recommended # instead we will convert this to a function to perform lazy initialization return LuxLinear(() -> copy(A), () -> copy(B)) end # `B` is a parameter Lux.initialparameters(::AbstractRNG, layer::LuxLinear) = (B=layer.init_B(),) # `A` is a state Lux.initialstates(::AbstractRNG, layer::LuxLinear) = (A=layer.init_A(),) (l::LuxLinear)(x, ps, st) = st.A * ps.B * x, st ``` ```julia [Flux] using Flux, Random, NNlib, Zygote, Optimisers struct FluxLinear A B end # `A` is not trainable Optimisers.trainable(f::FluxLinear) = (B=f.B,) # Needed so that both `A` and `B` can be transferred between devices Flux.@functor FluxLinear (l::FluxLinear)(x) = l.A * l.B * x ``` ::: Now let us run the model. :::code-group ```julia{2,5,7,9} [Lux] rng = Random.default_rng() model = LuxLinear(randn(rng, 2, 4), randn(rng, 4, 2)) x = randn(rng, 2, 1) ps, st = Lux.setup(rng, model) model(x, ps, st) gradient(ps -> sum(first(model(x, ps, st))), ps) ``` ```julia [Flux] rng = Random.default_rng() model = FluxLinear(randn(rng, 2, 4), randn(rng, 4, 2)) x = randn(rng, 2, 1) model(x) gradient(model -> sum(model(x)), model) ``` ::: To reiterate some important points: * Don't store mutables like Arrays inside a Lux Layer. * Parameters and States should be constructured inside the respective `initial*` functions. ## Certain Important Implementation Details {#Certain-Important-Implementation-Details} ### Training/Inference Mode {#Training/Inference-Mode} Flux supports a mode called `:auto` which automatically decides if the user is training the model or running inference. This is the default mode for `Flux.BatchNorm`, `Flux.GroupNorm`, `Flux.Dropout`, etc. Lux doesn't support this mode (specifically to keep code simple and do exactly what the user wants), hence our default mode is `training`. This can be changed using `Lux.testmode`. ## Can we still use Flux Layers? {#Can-we-still-use-Flux-Layers?} If you have `Flux` loaded in your code, you can use the function [`FromFluxAdaptor`](/api/Lux/interop#Lux.FromFluxAdaptor) to automatically convert your model to `Lux`. Note that in case a native Lux counterpart isn't available, we fallback to using `Optimisers.destructure`. --- --- url: /dev/manual/flux_lux_interop.md --- # Supporting Both Flux and Lux {#flux-lux-interop} A common question for package maintainers is: "How can I support both Flux and Lux in my package?" This guide provides a comprehensive approach to maintaining compatibility with both frameworks while minimizing code duplication and dependency overhead. ## The Core Strategy {#The-Core-Strategy} The recommended approach is to: 1. **Define your core layers using LuxCore**: Use `LuxCore.jl` as your primary interface since it's a lighter dependency than full `Lux.jl` 2. **Construct a StatefulLuxLayer**: Wrap the layer in a [`StatefulLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer) to provide a Flux-style interface This strategy allows users to choose their preferred framework while keeping your package's core functionality framework-agnostic. ## Implementation Pattern {#Implementation-Pattern} ### 1. Core Layer Definition {#1.-Core-Layer-Definition} First, define your layer using the LuxCore interface: ```julia using LuxCore, Random struct MyCustomLuxLayer{F} <: AbstractLuxLayer # Layer configuration (no mutable state!) feature_dim::Int output_dim::Int activation_fn::F end function MyCustomLuxLayer(feature_dim::Int, output_dim::Int; activation=identity) return MyCustomLuxLayer(feature_dim, output_dim, activation) end # Define the Lux interface function LuxCore.initialparameters(rng::AbstractRNG, layer::MyCustomLuxLayer) return ( weight = randn(rng, Float32, layer.output_dim, layer.feature_dim), bias = zeros(Float32, layer.output_dim) ) end function (layer::MyCustomLuxLayer)(x, ps, st) y = ps.weight * x .+ ps.bias y = layer.activation_fn.(y) return y, st end ``` ### 2. Wrap the layer in a StatefulLuxLayer {#2.-Wrap-the-layer-in-a-StatefulLuxLayer} [`StatefulLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer) is a convenience wrapper over Lux layers which stores the parameters and states (and handles updating the state internally). This layer is also compatible with `Flux.jl`. ## Usage Examples {#Usage-Examples} ### Using the Lux Interface {#Using-the-Lux-Interface} ```julia using LuxCore, Random, Flux # Create layer and setup rng = Random.default_rng() layer = MyCustomLuxLayer(4, 2; activation=tanh) ps, st = LuxCore.setup(rng, layer) # Forward pass x = randn(Float32, 4, 32) # batch of 32 samples y, st_new = layer(x, ps, st) ``` ```ansi (Float32[0.8389038 0.9993076 … 0.53918076 0.99998707; 0.9966145 -0.11863235 … 0.9981629 0.058118183], NamedTuple()) ``` ### Using the Flux Interface {#Using-the-Flux-Interface} ```julia using Flux, LuxCore, Random # Create Flux-style layer model = MyCustomLuxLayer(4, 2; activation=tanh) ps, st = LuxCore.setup(Random.default_rng(), model) flux_model = LuxCore.StatefulLuxLayer(model, ps, st) # Use like any Flux layer x = randn(Float32, 4, 32) y_target = randn(Float32, 2, 32) y = flux_model(x) # Works with Flux training using Optimisers opt = Adam(0.01) opt_state = Optimisers.setup(opt, flux_model) # Training step loss_fn(m, x, y_target) = Flux.mse(m(x), y_target) loss, grads = Flux.withgradient(loss_fn, flux_model, x, y_target) opt_state, flux_model = Optimisers.update(opt_state, flux_model, grads[1]) ``` ```ansi ((model = (), ps = (weight = Leaf(Adam(eta=0.01, beta=(0.9, 0.999), epsilon=1.0e-8), (Float32[-0.00738183 -0.00600417 -0.0319064 -0.000142127; 0.00584962 2.74381f-5 -0.0171094 -0.00460663], Float32[5.44906f-6 3.60496f-6 0.0001018 2.01998f-9; 3.42176f-6 7.52841f-11 2.92727f-5 2.12208f-6], (0.81, 0.998001))), bias = Leaf(Adam(eta=0.01, beta=(0.9, 0.999), epsilon=1.0e-8), (Float32[-0.0187933, -0.00834907], Float32[3.53183f-5, 6.9706f-6], (0.81, 0.998001)))), st = (), st_any = ()), StatefulLuxLayer{Val{true}, Main.MyCustomLuxLayer{typeof(tanh)}, @NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, @NamedTuple{}}(Main.MyCustomLuxLayer{typeof(tanh)}(4, 2, tanh), (weight = Float32[0.6957868 -0.40701553 -0.73492825 0.10898119; 1.4767253 0.08893849 -0.58276945 -0.5118015], bias = Float32[0.009999999, 0.009999999]), NamedTuple(), nothing, Val{true}())) ``` ## Best Practices {#Best-Practices} 1. **Use LuxCore for core definitions**: Depend on `LuxCore.jl` rather than full `Lux.jl` to minimize dependencies. 2. **Lazy loading**: Use package extensions to avoid loading Flux unless needed. ## Common Gotchas {#Common-Gotchas} 1. **Mutable state in layer structs**: Remember that Lux layers should not contain mutable state. Put mutable objects in the state, not the layer. 2. **Parameter sharing**: Be careful with parameter sharing when converting between interfaces. 3. **Extension loading**: Users need to load Flux explicitly to access the Flux interface, even if your package supports it. By following this pattern, you can provide excellent support for both Flux and Lux users while maintaining clean, maintainable code. --- --- url: /dev/manual/dispatch_custom_input.md --- # Dispatching on Custom Input Types {#Dispatching-on-Custom-Input-Types} ## Which function should participate in dispatch? {#Which-function-should-participate-in-dispatch?} * Defining a dispatch on `(::Layer)(x::MyInputType, ps, st::NamedTuple)` is inconvenient, since it requires the user to define a new method for every layer type. * `(::AbstractLuxLayer)(x::MyInputType, ps, st::NamedTuple)` doesn't work. * Instead, we need to define the dispatch on `Lux.apply(::AbstractLuxLayer, x::MyInputType, ps, st::NamedTuple)`. ## Concrete Example {#Concrete-Example} Consider [Neural ODEs](https://implicit-layers-tutorial.org/neural_odes/). In these models, often time we want to every iteration of the neural network to take the current time as input. Here, we won't go through implementing an entire Neural ODE model. Instead we will define a time dependent version of [`Chain`](/api/Lux/layers#Lux.Chain). ### Time-Dependent Chain Implementation {#Time-Dependent-Chain-Implementation} ```julia using Lux, Random struct TDChain{L <: NamedTuple} <: Lux.AbstractLuxWrapperLayer{:layers} layers::L end function (l::TDChain)((x, t)::Tuple, ps, st::NamedTuple) # Concatenate along the 2nd last dimension sz = ntuple(i -> i == ndims(x) - 1 ? 1 : size(x, i), ndims(x)) t_ = ones(eltype(x), sz) .* t # Needs to be modified for GPU for name in keys(l.layers) x, st_ = Lux.apply(getfield(l.layers, name), cat(x, t_; dims=ndims(x) - 1), getfield(ps, name), getfield(st, name)) st = merge(st, NamedTuple{(name,)}((st_,))) end return x, st end model = Chain(Dense(3, 4), TDChain((; d1=Dense(5, 4), d2=Dense(5, 4))), Dense(4, 1)) ``` ```ansi Chain( layer_1 = Dense(3 => 4), # 16 parameters layer_2 = TDChain( d(1-2) = Dense(5 => 4), # 48 (24 x 2) parameters ), layer_3 = Dense(4 => 1), # 5 parameters )  # Total: 69 parameters,  # plus 0 states. ``` ### Running the TDChain {#Running-the-TDChain} ```julia rng = MersenneTwister(0) ps, st = Lux.setup(rng, model) x = randn(rng, Float32, 3, 2) try model(x, ps, st) catch e Base.showerror(stdout, e) end ``` ```ansi MethodError: no method matching apply(::@NamedTuple{d1::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}, d2::Dense{typeof(identity), Int64, Int64, Nothing, Nothing, Static.True}}, ::Matrix{Float32}, ::@NamedTuple{d1::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}, d2::@NamedTuple{weight::Matrix{Float32}, bias::Vector{Float32}}}, ::@NamedTuple{d1::@NamedTuple{}, d2::@NamedTuple{}}) The function `apply` exists, but no method is defined for this combination of argument types. Closest candidates are:  apply(::AbstractLuxLayer, ::Any, ::Any, ::Any)  @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/LuxCore.jl:154  apply(::StatefulLuxLayer, ::Any, ::Any)  @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/stateful.jl:163  apply(::StatefulLuxLayer, ::Any)  @ LuxCore ~/work/Lux.jl/Lux.jl/lib/LuxCore/src/stateful.jl:163 ``` ### Writing the Correct Dispatch Rules {#Writing-the-Correct-Dispatch-Rules} * Create a Custom Layer storing the time. ```julia struct ArrayAndTime{A <: AbstractArray, T <: Real} array::A time::T end ``` * Define the dispatch on `Lux.apply(::AbstractLuxLayer, x::ArrayAndTime, ps, st::NamedTuple)`. ```julia function Lux.apply(layer::Lux.AbstractLuxLayer, x::ArrayAndTime, ps, st::NamedTuple) y, st = layer(x.array, ps, st) return ArrayAndTime(y, x.time), st end function Lux.apply(layer::TDChain, x::ArrayAndTime, ps, st::NamedTuple) y, st = layer((x.array, x.time), ps, st) return ArrayAndTime(y, x.time), st end ``` * Run the model. ```julia xt = ArrayAndTime(x, 10.0f0) model(xt, ps, st)[1] ``` ```ansi Main.ArrayAndTime{Matrix{Float32}, Float32}(Float32[4.887438 5.5271416], 10.0f0) ``` ### Using the Same Input for Non-TD Models {#Using-the-Same-Input-for-Non-TD-Models} Writing proper dispatch means we can simply replace the `TDChain` with a `Chain` (of course with dimension corrections) and the pipeline still works. ```julia model = Chain(Dense(3, 4), Chain((; d1=Dense(4, 4), d2=Dense(4, 4))), Dense(4, 1)) ps, st = Lux.setup(rng, model) model(xt, ps, st)[1] ``` ```ansi Main.ArrayAndTime{Matrix{Float32}, Float32}(Float32[0.40721768 1.2363781], 10.0f0) ``` --- --- url: /dev/manual/preferences.md --- # Preferences for Lux.jl {#Preferences-for-Lux.jl} ::: tip How to set Preferences [PreferenceTools.jl](https://github.com/cjdoris/PreferenceTools.jl) provides an interactive way to set preferences. First run the following command: ```julia using PreferenceTools ``` Then in the pkg mode (press `]` in the REPL), run the following command: ```julia pkg> preference add Lux = pkg> preference add LuxLib = pkg> preference add LuxCore = ``` ::: Lux.jl relies on several preferences to make decision on how to run your code. Here is an exhaustive list of preferences that Lux.jl uses. ## Nested Automatic Differentiation {#Nested-Automatic-Differentiation} 1. `automatic_nested_ad_switching` - Set this to `false` to disable automatic switching of backends for nested automatic differentiation. See the manual section on [nested automatic differentiation](/manual/nested_autodiff#nested_autodiff) for more details. ## Training with Reactant {#Training-with-Reactant} 1. `precision_config` - Set this to `"auto"` to use select `HIGH` precision for CUDA and `DEFAULT` precision for other backends. The choice made here can change across versions of Lux.jl and is generally recommended for best performance. Alternatively, you can set this to `"default"` to use `DEFAULT` precision or `"high"` to use `HIGH` precision or `"highest"` to use `HIGHEST` precision. ## GPU-Aware MPI Support {#gpu-aware-mpi-preferences} If you are using a custom MPI build that supports CUDA or ROCM, you can use the following preferences with [Preferences.jl](https://github.com/JuliaPackaging/Preferences.jl): 1. `cuda_aware_mpi` - Set this to `true` if your MPI build is CUDA aware. 2. `rocm_aware_mpi` - Set this to `true` if your MPI build is ROCM aware. By default, both of these preferences are set to `false`. ## GPU Backend Selection {#GPU-Backend-Selection} 1. `gpu_backend` - Set this to bypass the automatic backend selection and use a specific gpu backend. Valid options are "cuda", "rocm", "metal", and "oneapi". This preference needs to be set for `MLDataDevices` package. It is recommended to use [`MLDataDevices.gpu_backend!`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_backend!) to set this preference. ## Automatic Eltype Conversion {#automatic-eltypes-preference} 1. `eltype_mismatch_handling` - Preference controlling what happens when layers get different eltypes as input. See the documentation on [`match_eltype`](/api/Lux/utilities#Lux.match_eltype) for more details. ## Dispatch Doctor {#dispatch-doctor-preference} 1. `instability_check` - Preference controlling the dispatch doctor. See the documentation on [`Lux.set_dispatch_doctor_preferences!`](/api/Lux/utilities#Lux.set_dispatch_doctor_preferences!) for more details. The preferences need to be set for `LuxCore` and `LuxLib` packages. Both of them default to `disable`. * Setting the `LuxCore` preference sets the check at the level of `LuxCore.apply`. This essentially activates the dispatch doctor for all Lux layers. * Setting the `LuxLib` preference sets the check at the level of functional layer of Lux, for example, [`fused_dense_bias_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.fused_dense_bias_activation). These functions are supposed to be type stable for common input types and can be used to guarantee type stability. ## Disabling Loop Vectorization / Octavian {#disable\_loop\_vectorization} `LoopVectorization.jl` and `Octavian.jl` are optional dependencies that are used to accelerate certain CPU operations. However, these packages are tightly coupled with julia and might not work with all julia versions and systems. If these packages are loaded in any form LuxLib will use the optimized versions of the functions. But it might be desirable to disable these packages and use the default implementations instead. This can be done by setting the `disable_loop_vectorization` preference to `true` for `LuxLib`. --- --- url: /dev/manual/distributed_utils.md --- # Distributed Data Parallel Training {#Distributed-Data-Parallel-Training} ::: tip Tip For a fully functional example, see the [ImageNet Training Example](https://github.com/LuxDL/Lux.jl/tree/main/examples/ImageNet). ::: DDP Training using `Lux.DistributedUtils` is a spiritual successor to [FluxMPI.jl](https://github.com/avik-pal/FluxMPI.jl), but has some key differences. ## Guide to Integrating DistributedUtils into your code {#Guide-to-Integrating-DistributedUtils-into-your-code} * Initialize the respective backend with [`DistributedUtils.initialize`](/api/Lux/distributed_utils#Lux.DistributedUtils.initialize), by passing in a backend type. It is important that you pass in the type, i.e. `NCCLBackend` and not the object `NCCLBackend()`. ```julia DistributedUtils.initialize(NCCLBackend) ``` * Obtain the backend via [`DistributedUtils.get_distributed_backend`](/api/Lux/distributed_utils#Lux.DistributedUtils.get_distributed_backend) by passing in the type of the backend (same note as last point applies here again). ```julia backend = DistributedUtils.get_distributed_backend(NCCLBackend) ``` It is important that you use this function instead of directly constructing the backend, since there are certain internal states that need to be synchronized. * Next synchronize the parameters and states of the model. This is done by calling [`DistributedUtils.synchronize!!`](/api/Lux/distributed_utils#Lux.DistributedUtils.synchronize!!) with the backend and the respective input. ```julia ps = DistributedUtils.synchronize!!(backend, ps) st = DistributedUtils.synchronize!!(backend, st) ``` * To split the data uniformly across the processes use [`DistributedUtils.DistributedDataContainer`](/api/Lux/distributed_utils#Lux.DistributedUtils.DistributedDataContainer). Alternatively, one can manually split the data. For the provided container to work [`MLUtils.jl`](https://github.com/JuliaML/MLUtils.jl) must be installed and loaded. ```julia data = DistributedUtils.DistributedDataContainer(backend, data) ``` * Wrap the optimizer in [`DistributedUtils.DistributedOptimizer`](/api/Lux/distributed_utils#Lux.DistributedUtils.DistributedOptimizer) to ensure that the optimizer is correctly synchronized across all processes before parameter updates. After initializing the state of the optimizer, synchronize the state across all processes. ```julia opt = DistributedUtils.DistributedOptimizer(backend, opt) opt_state = Optimisers.setup(opt, ps) opt_state = DistributedUtils.synchronize!!(backend, opt_state) ``` * Finally change all logging and serialization code to trigger on `local_rank(backend) == 0`. This ensures that only the master process logs and serializes the model. ## Migration Guide from `FluxMPI.jl` {#Migration-Guide-from-FluxMPI.jl} Let's compare the changes we need to make wrt the [FluxMPI.jl integration guide](https://avik-pal.github.io/FluxMPI.jl/dev/guide/). 1. `FluxMPI.Init` is now [`DistributedUtils.initialize`](/api/Lux/distributed_utils#Lux.DistributedUtils.initialize). 2. `FluxMPI.synchronize!(x)` needs to be changed to `x_new = DistributedUtils.synchronize!!(backend, x)`. 3. [`DistributedUtils.DistributedDataContainer`](/api/Lux/distributed_utils#Lux.DistributedUtils.DistributedDataContainer), [`DistributedUtils.local_rank`](/api/Lux/distributed_utils#Lux.DistributedUtils.local_rank), and [`DistributedUtils.DistributedOptimizer`](/api/Lux/distributed_utils#Lux.DistributedUtils.DistributedOptimizer) need `backend` as the first input. And that's pretty much it! ### Removed Functionality {#Removed-Functionality} 1. `FluxMPI.allreduce_gradients` no longer exists. Previously this was needed when CUDA communication was flaky, with `NCCL.jl` this is no longer the case. 2. `FluxMPIFluxModel` has been removed. `DistributedUtils` no longer works with `Flux`. ### Key Differences {#Key-Differences} 1. `FluxMPI.synchronize!` is now `DistributedUtils.synchronize!!` to highlight the fact that some of the inputs are not updated in-place. 2. All of the functions now require a [communication backend](/api/Lux/distributed_utils#communication-backends) as input. 3. We don't automatically determine if the MPI Implementation is CUDA or ROCM aware. See [GPU-aware MPI](/manual/preferences#gpu-aware-mpi-preferences) for more information. 4. Older (now non-existent) `Lux.gpu` implementations used to "just work" with `FluxMPI.jl`. We expect [`gpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_device) to continue working as expected, however, we recommend using [`gpu_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.gpu_device) after calling [`DistributedUtils.initialize`](/api/Lux/distributed_utils#Lux.DistributedUtils.initialize) to avoid any mismatch between the device set via `DistributedUtils` and the device stores in `CUDADevice` or `AMDGPUDevice`. ## Known Shortcomings {#Known-Shortcomings} 1. Currently we don't run tests with CUDA or ROCM aware MPI, use those features at your own risk. We are working on adding tests for these features. 2. Native AMDGPU.jl support is experimental and causes deadlocks in certain situations. **For AMD GPUs, we strongly recommend using Reactant instead of native AMDGPU.jl for distributed training.** If you have a minimal reproducer for AMDGPU.jl issues, please open an issue. --- --- url: /dev/api/Lux/layers.md --- # Built-In Layers {#Built-In-Layers} ## Containers {#Containers} ```julia BranchLayer(layers...; fusion=nothing) BranchLayer(; fusion=nothing, name=nothing, layers...) ``` Takes an input `x` and passes it through all the `layers` and returns a tuple of the outputs. If `fusion` is provided, applies fusion to the tuple of outputs. **Arguments** * Layers can be specified in two formats: * A list of `N` Lux layers * Specified as `N` keyword arguments. **Keyword Arguments** * `fusion`: An optional layer or function to apply to the tuple of outputs. If `fusion = nothing`, returns the tuple as-is (default behavior). If `fusion` is provided, returns `fusion((layer_1(x), layer_2(x), ..., layer_N(x)))`. **Extended Help** **Inputs** * `x`: Will be directly passed to each of the `layers` **Returns** * If `fusion = nothing`: Tuple `(layer_1(x), layer_2(x), ..., layer_N(x))` (naming changes if using the kwargs API) * If `fusion` is provided: `fusion((layer_1(x), layer_2(x), ..., layer_N(x)))` * Updated state of the `layers` (and `fusion` if it's a layer) **Parameters** * Parameters of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) * If `fusion` is an AbstractLuxLayer, parameters include both `layers` and `fusion` **States** * States of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) * If `fusion` is an AbstractLuxLayer, states include both `layers` and `fusion` ::: tip Comparison with Parallel This is slightly different from [`Parallel(nothing, layers...)`](/api/Lux/layers#Lux.Parallel) * If the input is a tuple, `Parallel` will pass each element individually to each layer. * `BranchLayer` essentially assumes 1 input comes in and is branched out into `N` outputs. ::: **Example** An easy way to replicate an input to an NTuple is to do ```julia julia> BranchLayer(NoOpLayer(), NoOpLayer(), NoOpLayer()) BranchLayer( layer_(1-3) = NoOpLayer(), ) # Total: 0 parameters, # plus 0 states. ``` source ```julia Chain(layers...; name=nothing) Chain(; layers..., name=nothing) ``` Collects multiple layers / functions to be called in sequence on a given input. **Arguments** * Layers can be specified in two formats: * A list of `N` Lux layers * Specified as `N` keyword arguments. **Extended Help** **Inputs** Input `x` is passed sequentially to each layer, and must conform to the input requirements of the internal layers. **Returns** * Output after sequentially applying all the layers to `x` * Updated model states **Parameters** * Parameters of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) **States** * States of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) **Miscellaneous Properties** * Allows indexing and field access syntax. We can access the `i`th layer by `m[i]` or `m.layer_i`. We can also index using ranges or arrays. **Example** ```julia julia> Chain(Dense(2, 3, relu), BatchNorm(3), Dense(3, 2)) Chain( layer_1 = Dense(2 => 3, relu), # 9 parameters layer_2 = BatchNorm(3, affine=true, track_stats=true), # 6 parameters, plus 7 non-trainable layer_3 = Dense(3 => 2), # 8 parameters ) # Total: 23 parameters, # plus 7 states. julia> Chain(Dense(2, 3, relu), BatchNorm(3), Dense(3, 2); name="MyFancyChain") MyFancyChain( layer_1 = Dense(2 => 3, relu), # 9 parameters layer_2 = BatchNorm(3, affine=true, track_stats=true), # 6 parameters, plus 7 non-trainable layer_3 = Dense(3 => 2), # 8 parameters ) # Total: 23 parameters, # plus 7 states. ``` source ```julia PairwiseFusion(connection, layers...; name=nothing) PairwiseFusion(connection; name=nothing, layers...) PairwiseFusion(; connection, layers..., name=nothing) ``` ```julia x1 → layer1 → y1 ↘ connection → layer2 → y2 ↘ x2 ↗ connection → y3 x3 ↗ ``` **Arguments** * `connection`: Takes 2 inputs and combines them * `layers`: `AbstractLuxLayer`s. Layers can be specified in two formats: * A list of `N` Lux layers * Specified as `N` keyword arguments. **Extended Help** **Inputs** Layer behaves differently based on input type: 1. If the input `x` is a tuple of length `N + 1`, then the `layers` must be a tuple of length `N`. The computation is as follows ```julia y = x[1] for i in 1:N y = connection(x[i + 1], layers[i](y)) end ``` 2. Any other kind of input ```julia y = x for i in 1:N y = connection(x, layers[i](y)) end ``` **Returns** * See Inputs section for how the return value is computed * Updated model state for all the contained layers **Parameters** * Parameters of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) **States** * States of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) source ```julia Parallel(connection, layers...; name=nothing) Parallel(connection; name=nothing, layers...) Parallel(; connection, layers..., name=nothing) ``` Create a layer which passes an input to each path in `layers`, before reducing the output with `connection`. **Arguments** * `connection`: An `N`-argument function that is called after passing the input through each layer, OR an AbstractLuxLayer that takes a tuple of `N` inputs. If `connection = nothing`, we return a tuple: `Parallel(nothing, f, g)(x, y) = (f(x), g(y))` * Layers can be specified in two formats: * A list of `N` Lux layers * Specified as `N` keyword arguments. **Extended Help** **Inputs** * `x`: If `x` is not a tuple, then return is computed as `connection([l(x) for l in layers]...)`. Else one is passed to each layer, thus `Parallel(+, f, g)(x, y) = f(x) + g(y)`. **Returns** * See the Inputs section for how the output is computed * Updated state of the `layers` (and `connection` if it's a layer) **Parameters** * Parameters of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) * If `connection` is an AbstractLuxLayer, parameters include both `layers` and `connection` **States** * States of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) * If `connection` is an AbstractLuxLayer, states include both `layers` and `connection` See also [`SkipConnection`](/api/Lux/layers#Lux.SkipConnection) which is `Parallel` with one identity. **Example** ```julia julia> model = Parallel(nothing, Dense(2, 1), Dense(2, 1)) Parallel( layer_(1-2) = Dense(2 => 1), # 6 (3 x 2) parameters ) # Total: 6 parameters, # plus 0 states. julia> using Random; rng = Random.seed!(123); ps, st = Lux.setup(rng, model); x1 = randn(rng, Float32, 2); x2 = randn(rng, Float32, 2); julia> size.(first(model((x1, x2), ps, st))) ((1,), (1,)) ``` source ```julia SkipConnection(layers, connection; name=nothing) SkipConnection(; layers, connection, name=nothing) ``` Create a skip connection which consists of a layer or [`Chain`](/api/Lux/layers#Lux.Chain) of consecutive layers and a shortcut connection linking the block's input to the output through a user-supplied 2-argument callable. The first argument to the callable will be propagated through the given `layer` while the second is the unchanged, "skipped" input. The simplest "ResNet"-type connection is just `SkipConnection(layer, +)`. **Arguments** * `layer`: Layer or `Chain` of layers to be applied to the input * `connection`: * A 2-argument function that takes `layer(input)` and the input OR * An AbstractLuxLayer that takes `(layer(input), input)` as input **Extended Help** **Inputs** * `x`: Will be passed directly to `layer` **Returns** * Output of `connection(layer(input), input)` * Updated state of `layer` **Parameters** * Parameters of `layer` OR * If `connection` is an AbstractLuxLayer, then NamedTuple with fields `:layers` and `:connection` **States** * States of `layer` OR * If `connection` is an AbstractLuxLayer, then NamedTuple with fields `:layers` and `:connection` See [`Parallel`](/api/Lux/layers#Lux.Parallel) for a more general implementation. source ```julia RepeatedLayer(model; repeats::Val = Val(10), input_injection::Val = Val(false)) ``` Iteratively applies `model` for `repeats` number of times. The initial input is passed into the model repeatedly if `input_injection = Val(true)`. This layer unrolls the computation, however, semantically this is same as: * `input_injection = Val(false)` ```julia res = x for i in 1:repeats res, st = model(res, ps, st) end ``` * `input_injection = Val(true)` ```julia res = x for i in 1:repeats res, st = model((res, x), ps, st) end ``` It is expected that `repeats` will be a reasonable number below `20`, beyond that compile times for gradients might be unreasonably high. **Arguments** * `model` must be an `AbstractLuxLayer` **Keyword Arguments** * `repeats`: Number of times to apply the model * `input_injection`: If `true`, then the input is passed to the model along with the output **Extended Help** **Inputs** * `x`: Input as described above **Returns** * Output is computed by as described above * Updated state of the `model` **Parameters** * Parameters of `model` **States** * State of `model` source ```julia AlternatePrecision{T}(layer) AlternatePrecision(::Type{T}, layer) ``` This layer is used to convert the input to a different precision (`T`), execute the layer, and then convert the output back to the original precision. **Arguments** * `T`: The eltype of the input to the layer * `layer`: The layer to execute **Inputs** * `x`: AbstractArray **Returns** * `y`: Output of the layer * State of the output source ## Convolutional Layers {#Convolutional-Layers} ```julia Conv(k::NTuple{N,Integer}, (in_chs => out_chs)::Pair{<:Integer,<:Integer}, activation=identity; init_weight=nothing, init_bias=nothing, stride=1, pad=0, dilation=1, groups=1, use_bias=True(), cross_correlation=False()) ``` Standard convolutional layer. ::: tip Conv 2D Image data should be stored in WHCN order (width, height, channels, batch). In other words, a `100 x 100` RGB image would be a `100 x 100 x 3 x 1` array, and a batch of 50 would be a `100 x 100 x 3 x 50` array. This has `N = 2` spatial dimensions, and needs a kernel size like `(5, 5)`, a 2-tuple of integers. To take convolutions along `N` feature dimensions, this layer expects as input an array with `ndims(x) == N + 2`, where `size(x, N + 1) == in_chs` is the number of input channels, and `size(x, ndims(x))` is the number of observations in a batch. ::: ::: warning Warning Frameworks like [`Pytorch`](https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html#torch.nn.Conv2d) perform cross-correlation in their convolution layers. Pass `cross_correlation=true` to use cross-correlation instead. ::: **Arguments** * `k`: Tuple of integers specifying the size of the convolutional kernel. Eg, for 2D convolutions `length(k) == 2` * `in_chs`: Number of input channels * `out_chs`: Number of input and output channels * `activation`: Activation Function **Extended Help** **Keyword Arguments** * `init_weight`: Controls the initialization of the weight parameter. If `nothing`, then we use [`kaiming_uniform`](/api/Building_Blocks/WeightInitializers#WeightInitializers.kaiming_uniform) with gain computed on the basis of the activation function (taken from Pytorch [`nn.init.calculate_gain`](https://pytorch.org/docs/stable/nn.init.html#torch.nn.init.calculate_gain)). * `init_bias`: Controls the initialization of the bias parameter. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(fan_in))`. * `stride`: Should each be either single integer, or a tuple with `N` integers * `dilation`: Should each be either single integer, or a tuple with `N` integers * `pad`: Specifies the number of elements added to the borders of the data array. It can be * a single integer for equal padding all around, * a tuple of `N` integers, to apply the same padding at begin/end of each spatial dimension, * a tuple of `2*N` integers, for asymmetric padding, or * the singleton `SamePad()`, to calculate padding such that `size(output,d) == size(x,d) / stride` (possibly rounded) for each spatial dimension. * Periodic padding can achieved by pre-empting the layer with a `WrappedFunction(x -> NNlib.pad_circular(x, N_pad; dims=pad_dims))` * `groups`: Expected to be an `Int`. It specifies the number of groups to divide a convolution into (set `groups = in_chs` for Depthwise Convolutions). `in_chs` and `out_chs` must be divisible by `groups`. * `use_bias`: Trainable bias can be disabled entirely by setting this to `false`. * `cross_correlation`: If `true`, perform cross-correlation instead of convolution. Prior to `v1`, Lux used to have a `CrossCor` layer which performed cross-correlation. This was removed in `v1` in favor of `Conv` with `cross_correlation=true`. **Inputs** * `x`: Data satisfying `ndims(x) == N + 2 && size(x, N - 1) == in_chs`, i.e. `size(x) = (I_N, ..., I_1, C_in, N)` **Returns** * Output of the convolution `y` of size `(O_N, ..., O_1, C_out, N)` where $$O\_i = \left\lfloor\frac{I\_i + p\_i + p\_{(i + N) % |p|} - d\_i \times (k\_i - 1)}{s\_i} + 1\right\rfloor$$ * Empty `NamedTuple()` **Parameters** * `weight`: Convolution kernel * `bias`: Bias (present if `use_bias=true`) source ```julia ConvTranspose(k::NTuple{N,Integer}, (in_chs => out_chs)::Pair{<:Integer,<:Integer}, activation=identity; init_weight=glorot_uniform, init_bias=zeros32, stride=1, pad=0, outpad=0, dilation=1, groups=1, use_bias=True(), cross_correlation=False()) ``` Standard convolutional transpose layer. **Arguments** * `k`: Tuple of integers specifying the size of the convolutional kernel. Eg, for 2D convolutions `length(k) == 2` * `in_chs`: Number of input channels * `out_chs`: Number of input and output channels * `activation`: Activation Function **Keyword Arguments** * `init_weight`: Controls the initialization of the weight parameter. If `nothing`, then we use [`kaiming_uniform`](/api/Building_Blocks/WeightInitializers#WeightInitializers.kaiming_uniform) with gain computed on the basis of the activation function (taken from Pytorch [`nn.init.calculate_gain`](https://pytorch.org/docs/stable/nn.init.html#torch.nn.init.calculate_gain)). * `init_bias`: Controls the initialization of the bias parameter. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(fan_in))`. * `stride`: Should each be either single integer, or a tuple with `N` integers * `dilation`: Should each be either single integer, or a tuple with `N` integers * `pad`: Specifies the number of elements added to the borders of the data array. It can be * a single integer for equal padding all around, * a tuple of `N` integers, to apply the same padding at begin/end of each spatial dimension, * a tuple of `2*N` integers, for asymmetric padding, or * the singleton `SamePad()`, to calculate padding such that `size(output,d) == size(x,d) * stride` (possibly rounded) for each spatial dimension. * `groups`: Expected to be an `Int`. It specifies the number of groups to divide a convolution into (set `groups = in_chs` for Depthwise Convolutions). `in_chs` and `out_chs` must be divisible by `groups`. * `use_bias`: Trainable bias can be disabled entirely by setting this to `false`. * `cross_correlation`: If `true`, perform transposed cross-correlation instead of transposed convolution. * `outpad`: To converse [`Conv`](/api/Lux/layers#Lux.Conv) inversability when `stride > 1`, `outpad` can be used to increase the size of the output in the desired dimensions. Whereas `pad` is used to zero-pad the input, `outpad` only affects the output shape. **Extended Help** **Inputs** * `x`: Data satisfying `ndims(x) == N + 2 && size(x, N - 1) == in_chs`, i.e. `size(x) = (I_N, ..., I_1, C_in, N)` **Returns** * Output of the convolution transpose `y` of size `(O_N, ..., O_1, C_out, N)` where * Empty `NamedTuple()` **Parameters** * `weight`: Convolution Transpose kernel * `bias`: Bias (present if `use_bias=true`) source ## Dropout Layers {#Dropout-Layers} ```julia AlphaDropout(p::Real) ``` AlphaDropout layer. **Arguments** * `p`: Probability of Dropout * if `p = 0` then [`NoOpLayer`](/api/Lux/layers#Lux.NoOpLayer) is returned. * if `p = 1` then `WrappedLayer(Base.Fix1(broadcast, zero))` is returned. **Inputs** * `x`: Must be an AbstractArray **Returns** * `x` with dropout mask applied if `training=Val(true)` else just `x` * State with updated `rng` **States** * `rng`: Pseudo Random Number Generator * `training`: Used to check if training/inference mode Call [`Lux.testmode`](/api/Building_Blocks/LuxCore#LuxCore.testmode) to switch to test mode. See also [`Lux.Dropout`](/api/Lux/layers#Lux.Dropout), [`VariationalHiddenDropout`](/api/Lux/layers#Lux.VariationalHiddenDropout) source ```julia Dropout(p; dims=:) ``` Dropout layer. **Arguments** * `p`: Probability of Dropout (if `p = 0` then [`NoOpLayer`](/api/Lux/layers#Lux.NoOpLayer) is returned) **Keyword Arguments** * To apply dropout along certain dimension(s), specify the `dims` keyword. e.g. `Dropout(p; dims = (3,4))` will randomly zero out entire channels on WHCN input (also called 2D dropout). **Inputs** * `x`: Must be an AbstractArray **Returns** * `x` with dropout mask applied if `training=Val(true)` else just `x` * State with updated `rng` **States** * `rng`: Pseudo Random Number Generator * `training`: Used to check if training/inference mode Call [`Lux.testmode`](/api/Building_Blocks/LuxCore#LuxCore.testmode) to switch to test mode. See also [`AlphaDropout`](/api/Lux/layers#Lux.AlphaDropout), [`VariationalHiddenDropout`](/api/Lux/layers#Lux.VariationalHiddenDropout) source ```julia VariationalHiddenDropout(p; dims=:) ``` VariationalHiddenDropout layer. The only difference from Dropout is that the `mask` is retained until [`Lux.update_state(l, :update_mask, Val(true))`](/api/Building_Blocks/LuxCore#LuxCore.update_state) is called. **Arguments** * `p`: Probability of Dropout (if `p = 0` then [`NoOpLayer`](/api/Lux/layers#Lux.NoOpLayer) is returned) **Keyword Arguments** * To apply dropout along certain dimension(s), specify the `dims` keyword. e.g. `VariationalHiddenDropout(p; dims = 3)` will randomly zero out entire channels on WHCN input (also called 2D dropout). **Inputs** * `x`: Must be an AbstractArray **Returns** * `x` with dropout mask applied if `training=Val(true)` else just `x` * State with updated `rng` **States** * `rng`: Pseudo Random Number Generator * `training`: Used to check if training/inference mode * `mask`: Dropout mask. Initilly set to nothing. After every run, contains the mask applied in that call * `update_mask`: Stores whether new mask needs to be generated in the current call Call [`Lux.testmode`](/api/Building_Blocks/LuxCore#LuxCore.testmode) to switch to test mode. See also [`AlphaDropout`](/api/Lux/layers#Lux.AlphaDropout), [`Lux.Dropout`](/api/Lux/layers#Lux.Dropout) source ## Pooling Layers {#Pooling-Layers} ```julia AdaptiveLPPool(output_size; p=2) ``` Adaptive LP Pooling layer. Calculates the necessary window size such that its output has `size(y)[1:N] == output_size`. **Arguments** * `output_size`: Size of the first `N` dimensions for the output ::: danger GPU Support This layer is currently only supported on CPU. ::: **Inputs** * `x`: Expects as input an array with `ndims(x) == N + 2`, i.e. channel and batch dimensions, after the `N` feature dimensions, where `N = length(output_size)`. **Returns** * Output of size `(out..., C, N)` * Empty `NamedTuple()` source ```julia AdaptiveMaxPool(output_size) ``` Adaptive Max Pooling layer. Calculates the necessary window size such that its output has `size(y)[1:N] == output_size`. **Arguments** * `output_size`: Size of the first `N` dimensions for the output **Inputs** * `x`: Expects as input an array with `ndims(x) == N + 2`, i.e. channel and batch dimensions, after the `N` feature dimensions, where `N = length(output_size)`. **Returns** * Output of size `(out..., C, N)` * Empty `NamedTuple()` source ```julia AdaptiveMeanPool(output_size) ``` Adaptive Mean Pooling layer. Calculates the necessary window size such that its output has `size(y)[1:N] == output_size`. **Arguments** * `output_size`: Size of the first `N` dimensions for the output **Inputs** * `x`: Expects as input an array with `ndims(x) == N + 2`, i.e. channel and batch dimensions, after the `N` feature dimensions, where `N = length(output_size)`. **Returns** * Output of size `(out..., C, N)` * Empty `NamedTuple()` source ```julia GlobalLPPool(; p=2) ``` Global LP Pooling layer. Transforms `(w, h, c, b)`-shaped input into `(1, 1, c, b)`-shaped output, by performing mean pooling on the complete `(w, h)`-shaped feature maps. ::: danger GPU Support This layer is currently only supported on CPU. ::: **Inputs** * `x`: Data satisfying `ndims(x) > 2`, i.e. `size(x) = (I_N, ..., I_1, C, N)` **Returns** * Output of the pooling `y` of size `(1, ..., 1, C, N)` * Empty `NamedTuple()` source ```julia GlobalMaxPool() ``` Global Max Pooling layer. Transforms `(w, h, c, b)`-shaped input into `(1, 1, c, b)`-shaped output, by performing mean pooling on the complete `(w, h)`-shaped feature maps. **Inputs** * `x`: Data satisfying `ndims(x) > 2`, i.e. `size(x) = (I_N, ..., I_1, C, N)` **Returns** * Output of the pooling `y` of size `(1, ..., 1, C, N)` * Empty `NamedTuple()` source ```julia GlobalMeanPool() ``` Global Mean Pooling layer. Transforms `(w, h, c, b)`-shaped input into `(1, 1, c, b)`-shaped output, by performing mean pooling on the complete `(w, h)`-shaped feature maps. **Inputs** * `x`: Data satisfying `ndims(x) > 2`, i.e. `size(x) = (I_N, ..., I_1, C, N)` **Returns** * Output of the pooling `y` of size `(1, ..., 1, C, N)` * Empty `NamedTuple()` source ```julia LPPool(window; stride=window, pad=0, dilation=1, p=2) ``` LP Pooling layer, which replaces all pixels in a block of size `window` with the reduction operation: lp. **Arguments** * `window`: Tuple of integers specifying the size of the window. Eg, for 2D pooling `length(window) == 2` **Keyword Arguments** * `stride`: Should each be either single integer, or a tuple with `N` integers * `dilation`: Should each be either single integer, or a tuple with `N` integers * `pad`: Specifies the number of elements added to the borders of the data array. It can be * a single integer for equal padding all around, * a tuple of `N` integers, to apply the same padding at begin/end of each spatial dimension, * a tuple of `2*N` integers, for asymmetric padding, or * the singleton `SamePad()`, to calculate padding such that `size(output,d) == size(x,d) / stride` (possibly rounded) for each spatial dimension. ::: danger GPU Support This layer is currently only supported on CPU. ::: **Extended Help** **Inputs** * `x`: Data satisfying `ndims(x) == N + 2`, i.e. `size(x) = (I_N, ..., I_1, C, N)` **Returns** * Output of the pooling `y` of size `(O_N, ..., O_1, C, N)` where $$ O\_i = \left\lfloor\frac{I\_i + p\_i + p\_{(i + N) % |p|} - d\_i \times (k\_i - 1)}{s\_i} + 1\right\rfloor$$ * Empty `NamedTuple()` source ```julia MaxPool(window; stride=window, pad=0, dilation=1) ``` Max Pooling layer, which replaces all pixels in a block of size `window` with the reduction operation: max. **Arguments** * `window`: Tuple of integers specifying the size of the window. Eg, for 2D pooling `length(window) == 2` **Keyword Arguments** * `stride`: Should each be either single integer, or a tuple with `N` integers * `dilation`: Should each be either single integer, or a tuple with `N` integers * `pad`: Specifies the number of elements added to the borders of the data array. It can be * a single integer for equal padding all around, * a tuple of `N` integers, to apply the same padding at begin/end of each spatial dimension, * a tuple of `2*N` integers, for asymmetric padding, or * the singleton `SamePad()`, to calculate padding such that `size(output,d) == size(x,d) / stride` (possibly rounded) for each spatial dimension. **Extended Help** **Inputs** * `x`: Data satisfying `ndims(x) == N + 2`, i.e. `size(x) = (I_N, ..., I_1, C, N)` **Returns** * Output of the pooling `y` of size `(O_N, ..., O_1, C, N)` where $$ O\_i = \left\lfloor\frac{I\_i + p\_i + p\_{(i + N) % |p|} - d\_i \times (k\_i - 1)}{s\_i} + 1\right\rfloor$$ * Empty `NamedTuple()` source ```julia MeanPool(window; stride=window, pad=0, dilation=1) ``` Mean Pooling layer, which replaces all pixels in a block of size `window` with the reduction operation: mean. **Arguments** * `window`: Tuple of integers specifying the size of the window. Eg, for 2D pooling `length(window) == 2` **Keyword Arguments** * `stride`: Should each be either single integer, or a tuple with `N` integers * `dilation`: Should each be either single integer, or a tuple with `N` integers * `pad`: Specifies the number of elements added to the borders of the data array. It can be * a single integer for equal padding all around, * a tuple of `N` integers, to apply the same padding at begin/end of each spatial dimension, * a tuple of `2*N` integers, for asymmetric padding, or * the singleton `SamePad()`, to calculate padding such that `size(output,d) == size(x,d) / stride` (possibly rounded) for each spatial dimension. **Extended Help** **Inputs** * `x`: Data satisfying `ndims(x) == N + 2`, i.e. `size(x) = (I_N, ..., I_1, C, N)` **Returns** * Output of the pooling `y` of size `(O_N, ..., O_1, C, N)` where $$ O\_i = \left\lfloor\frac{I\_i + p\_i + p\_{(i + N) % |p|} - d\_i \times (k\_i - 1)}{s\_i} + 1\right\rfloor$$ * Empty `NamedTuple()` source ## Recurrent Layers {#Recurrent-Layers} ```julia GRUCell((in_dims, out_dims)::Pair{<:Int,<:Int}; use_bias=true, train_state::Bool=false, init_weight=glorot_uniform, init_recurrent_weight=init_weight, init_bias=nothing, init_state=zeros32) ``` Gated Recurrent Unit (GRU) Cell $$\begin{align} r &= \sigma(W\_{ir} \times x + b\_{ir} + W\_{hr} \times h\_{prev} + b\_{hr})\\ z &= \sigma(W\_{iz} \times x + b\_{iz} + W\_{hz} \times h\_{prev} + b\_{hz})\\ n &= \tanh(W\_{in} \times x + b\_{in} + r \cdot (W\_{hn} \times h\_{prev} + b\_{hn}))\\ h\_{new} &= (1 - z) \cdot n + z \cdot h\_{prev} \end{align}$$ **Arguments** * `in_dims`: Input Dimension * `out_dims`: Output (Hidden State) Dimension * `use_bias`: Set to false to deactivate bias * `train_state`: Trainable initial hidden state can be activated by setting this to `true` * `init_bias`: Initializer for bias. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3 element tuple. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_weight`: Initializer for weight. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3 element tuple. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_recurrent_weight`: Initializer for weight. Must be a tuple containing 3 functions. If a single value is passed, it is copied into a 3 element tuple. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_state`: Initializer for hidden state **Inputs** * Case 1a: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `false` - Creates a hidden state using `init_state` and proceeds to Case 2. * Case 1b: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `true` - Repeats `hidden_state` from parameters to match the shape of `x` and proceeds to Case 2. * Case 2: Tuple `(x, (h, ))` is provided, then the output and a tuple containing the updated hidden state is returned. **Returns** * Tuple containing * Output $h\_{new}$ of shape `(out_dims, batch_size)` * Tuple containing new hidden state $h\_{new}$ * Updated model state **Parameters** * `weight_ih`: Concatenated Weights to map from input space ${ W\_{ir}, W\_{iz}, W\_{in} }$. * `weight_hh`: Concatenated Weights to map from hidden space ${ W\_{hr}, W\_{hz}, W\_{hn} }$. * `bias_ih`: Concatenated Bias vector for the input space ${ b\_{ir}, b\_{iz}, b\_{in} }$ (not present if `use_bias=false`). * `bias_hh`: Concatenated Bias vector for the hidden space ${ b\_{hr}, b\_{hz}, b\_{hn} }$ (not present if `use_bias=false`). * `hidden_state`: Initial hidden state vector (not present if `train_state=false`) ${ b\_{hr}, b\_{hz}, b\_{hn} }$. **States** * `rng`: Controls the randomness (if any) in the initial state generation source ```julia LSTMCell(in_dims => out_dims; use_bias::Bool=true, train_state::Bool=false, train_memory::Bool=false, init_weight=nothing, init_recurrent_weight=init_weight, init_bias=nothing, init_state=zeros32, init_memory=zeros32) ``` Long Short-Term (LSTM) Cell $$\begin{align} i &= \sigma(W\_{ii} \times x + W\_{hi} \times h\_{prev} + b\_{i})\\ f &= \sigma(W\_{if} \times x + W\_{hf} \times h\_{prev} + b\_{f})\\ g &= \tanh(W\_{ig} \times x + W\_{hg} \times h\_{prev} + b\_{g})\\ o &= \sigma(W\_{io} \times x + W\_{ho} \times h\_{prev} + b\_{o})\\ c\_{new} &= f \cdot c\_{prev} + i \cdot g\\ h\_{new} &= o \cdot \tanh(c\_{new}) \end{align}$$ **Arguments** * `in_dims`: Input Dimension * `out_dims`: Output (Hidden State & Memory) Dimension * `use_bias`: Set to false to deactivate bias * `train_state`: Trainable initial hidden state can be activated by setting this to `true` * `train_memory`: Trainable initial memory can be activated by setting this to `true` * `init_bias`: Initializer for bias. Must be a tuple containing 4 functions. If a single value is passed, it is copied into a 4 element tuple. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_weight`: Initializer for weight. Must be a tuple containing 4 functions. If a single value is passed, it is copied into a 4 element tuple. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_recurrent_weight`: Initializer for recurrent weight. Must be a tuple containing 4 functions. If a single value is passed, it is copied into a 4 element tuple. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_state`: Initializer for hidden state * `init_memory`: Initializer for memory **Inputs** * Case 1a: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `false`, `train_memory` is set to `false` - Creates a hidden state using `init_state`, hidden memory using `init_memory` and proceeds to Case 2. * Case 1b: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `true`, `train_memory` is set to `false` - Repeats `hidden_state` vector from the parameters to match the shape of `x`, creates hidden memory using `init_memory` and proceeds to Case 2. * Case 1c: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `false`, `train_memory` is set to `true` - Creates a hidden state using `init_state`, repeats the memory vector from parameters to match the shape of `x` and proceeds to Case 2. * Case 1d: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `true`, `train_memory` is set to `true` - Repeats the hidden state and memory vectors from the parameters to match the shape of `x` and proceeds to Case 2. * Case 2: Tuple `(x, (h, c))` is provided, then the output and a tuple containing the updated hidden state and memory is returned. **Returns** * Tuple Containing * Output $h\_{new}$ of shape `(out_dims, batch_size)` * Tuple containing new hidden state $h\_{new}$ and new memory $c\_{new}$ * Updated model state **Parameters** * `weight_ih`: Concatenated Weights to map from input space ${ W\_{ii}, W\_{if}, W\_{ig}, W\_{io} }$. * `weight_hh`: Concatenated Weights to map from hidden space ${ W\_{hi}, W\_{hf}, W\_{hg}, W\_{ho} }$ * `bias_ih`: Bias vector for the input-hidden connection (not present if `use_bias=false`) * `bias_hh`: Concatenated Bias vector for the hidden-hidden connection (not present if `use_bias=false`) * `hidden_state`: Initial hidden state vector (not present if `train_state=false`) * `memory`: Initial memory vector (not present if `train_memory=false`) **States** * `rng`: Controls the randomness (if any) in the initial state generation source ```julia RNNCell(in_dims => out_dims, activation=tanh; use_bias=True(), train_state=False(), init_bias=nothing, init_weight=nothing, init_recurrent_weight=init_weight, init_state=zeros32) ``` An Elman RNNCell cell with `activation` (typically set to `tanh` or `relu`). $h\_{new} = activation(weight\_{ih} \times x + bias\_{ih} + weight\_{hh} \times h\_{prev} + bias\_{hh})$ **Arguments** * `in_dims`: Input Dimension * `out_dims`: Output (Hidden State) Dimension * `activation`: Activation function * `use_bias`: Set to false to deactivate bias * `train_state`: Trainable initial hidden state can be activated by setting this to `true` * `init_bias`: Initializer for bias. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_weight`: Initializer for weight. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_recurrent_weight`: Initializer for recurrent weight. If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(out_dims))`. * `init_state`: Initializer for hidden state **Inputs** * Case 1a: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `false` - Creates a hidden state using `init_state` and proceeds to Case 2. * Case 1b: Only a single input `x` of shape `(in_dims, batch_size)`, `train_state` is set to `true` - Repeats `hidden_state` from parameters to match the shape of `x` and proceeds to Case 2. * Case 2: Tuple `(x, (h, ))` is provided, then the output and a tuple containing the updated hidden state is returned. **Returns** * Tuple containing * Output $h\_{new}$ of shape `(out_dims, batch_size)` * Tuple containing new hidden state $h\_{new}$ * Updated model state **Parameters** * `weight_ih`: Maps the input to the hidden state. * `weight_hh`: Maps the hidden state to the hidden state. * `bias_ih`: Bias vector for the input-hidden connection (not present if `use_bias=false`) * `bias_hh`: Bias vector for the hidden-hidden connection (not present if `use_bias=false`) * `hidden_state`: Initial hidden state vector (not present if `train_state=false`) **States** * `rng`: Controls the randomness (if any) in the initial state generation source ```julia Recurrence( cell; ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex(), return_sequence::Bool=false, mincut::Bool=false, ) ``` Wraps a recurrent cell (like [`RNNCell`](/api/Lux/layers#Lux.RNNCell), [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell), [`GRUCell`](/api/Lux/layers#Lux.GRUCell)) to automatically operate over a sequence of inputs. ::: warning Relation to `Flux.Recur` This is completely distinct from `Flux.Recur`. It doesn't make the `cell` stateful, rather allows operating on an entire sequence of inputs at once. See [`StatefulRecurrentCell`](/api/Lux/layers#Lux.StatefulRecurrentCell) for functionality similar to `Flux.Recur`. ::: **Arguments** * `cell`: A recurrent cell. See [`RNNCell`](/api/Lux/layers#Lux.RNNCell), [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell), [`GRUCell`](/api/Lux/layers#Lux.GRUCell), for how the inputs/outputs of a recurrent cell must be structured. **Keyword Arguments** * `return_sequence`: If `true` returns the entire sequence of outputs, else returns only the last output. Defaults to `false`. * `ordering`: The ordering of the batch and time dimensions in the input. Defaults to `BatchLastIndex()`. Alternatively can be set to `TimeLastIndex()`. * `mincut`: If `true`, we will using mincut for the reverse mode differentiation. *(Only for Reactant)* **Extended Help** **Inputs** * If `x` is a * Tuple or Vector: Each element is fed to the `cell` sequentially. * Array (except a Vector): It is spliced along the penultimate dimension and each slice is fed to the `cell` sequentially. **Returns** * Output of the `cell` for the entire sequence. * Update state of the `cell`. ::: tip Tip Frameworks like Tensorflow have special implementation of [`StackedRNNCells`](https://www.tensorflow.org/api_docs/python/tf/keras/layers/StackedRNNCells) to handle sequentially composed RNN Cells. In Lux, one can simple stack multiple `Recurrence` blocks in a `Chain` to achieve the same. ```julia Chain( Recurrence(RNNCell(inputsize => latentsize); return_sequence=true), Recurrence(RNNCell(latentsize => latentsize); return_sequence=true), : x -> stack(x; dims=2) ) ``` For some discussion on this topic, see https://github.com/LuxDL/Lux.jl/issues/472. ::: source ```julia StatefulRecurrentCell(cell) ``` Wraps a recurrent cell (like [`RNNCell`](/api/Lux/layers#Lux.RNNCell), [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell), [`GRUCell`](/api/Lux/layers#Lux.GRUCell)) and makes it stateful. To avoid undefined behavior, once the processing of a single sequence of data is complete, update the state with `Lux.update_state(st, :carry, nothing)`. **Arguments** * `cell`: A recurrent cell. See [`RNNCell`](/api/Lux/layers#Lux.RNNCell), [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell), [`GRUCell`](/api/Lux/layers#Lux.GRUCell), for how the inputs/outputs of a recurrent cell must be structured. **Inputs** * Input to the `cell`. **Returns** * Output of the `cell` for the entire sequence. * Update state of the `cell` and updated `carry`. **States** * NamedTuple containing: * `cell`: Same as `cell`. * `carry`: The carry state of the `cell`. source ```julia BidirectionalRNN(cell::AbstractRecurrentCell, backward_cell::Union{AbstractRecurrentCell, Nothing}=nothing; merge_mode::Union{Function, Nothing}=vcat, ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex()) ``` Bidirectional RNN wrapper. **Arguments** * `cell`: A recurrent cell. See [`RNNCell`](/api/Lux/layers#Lux.RNNCell), [`LSTMCell`](/api/Lux/layers#Lux.LSTMCell), [`GRUCell`](/api/Lux/layers#Lux.GRUCell), for how the inputs/outputs of a recurrent cell must be structured. * `backward_cell`: A optional backward recurrent cell. If `backward_cell` is `nothing`, the rnn layer instance passed as the `cell` argument will be used to generate the backward layer automatically. `in_dims` of `backward_cell` should be consistent with `in_dims` of `cell` **Keyword Arguments** * `merge_mode`: Function by which outputs of the forward and backward RNNs will be combined. default value is `vcat`. If `nothing`, the outputs will not be combined. * `ordering`: The ordering of the batch and time dimensions in the input. Defaults to `BatchLastIndex()`. Alternatively can be set to `TimeLastIndex()`. **Extended Help** **Inputs** * If `x` is a * Tuple or Vector: Each element is fed to the `cell` sequentially. * Array (except a Vector): It is spliced along the penultimate dimension and each slice is fed to the `cell` sequentially. **Returns** * Merged output of the `cell` and `backward_cell` for the entire sequence. * Update state of the `cell` and `backward_cell`. **Parameters** * `NamedTuple` with `cell` and `backward_cell`. **States** * Same as `cell` and `backward_cell`. source ## Linear Layers {#Linear-Layers} ```julia Bilinear((in1_dims, in2_dims) => out, activation=identity; init_weight=nothing, init_bias=nothing, use_bias=True()) Bilinear(in12_dims => out, activation=identity; init_weight=nothing, init_bias=nothing, use_bias=True()) ``` Create a fully connected layer between two inputs and an output, and otherwise similar to [`Dense`](/api/Lux/layers#Lux.Dense). Its output, given vectors `x` & `y`, is another vector `z` with, for all `i in 1:out`: `z[i] = activation(x' * W[i, :, :] * y + bias[i])` If `x` and `y` are matrices, then each column of the output `z = B(x, y)` is of this form, with `B` the Bilinear layer. **Arguments** * `in1_dims`: number of input dimensions of `x` * `in2_dims`: number of input dimensions of `y` * `in12_dims`: If specified, then `in1_dims = in2_dims = in12_dims` * `out`: number of output dimensions * `activation`: activation function **Keyword Arguments** * `init_weight`: initializer for the weight matrix (`weight = init_weight(rng, out_dims, in1_dims, in2_dims)`). If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(in1_dims))`. * `init_bias`: initializer for the bias vector (ignored if `use_bias=false`). If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(in1_dims))`. * `use_bias`: Trainable bias can be disabled entirely by setting this to `false` **Input** * A 2-Tuple containing * `x` must be an AbstractArray with `size(x, 1) == in1_dims` * `y` must be an AbstractArray with `size(y, 1) == in2_dims` * If the input is an AbstractArray, then `x = y` **Returns** * AbstractArray with dimensions `(out_dims, size(x, 2))` * Empty `NamedTuple()` **Parameters** * `weight`: Weight Matrix of size `(out_dims, in1_dims, in2_dims)` * `bias`: Bias of size `(out_dims, 1)` (present if `use_bias=true`) source ```julia Dense(in_dims => out_dims, activation=identity; init_weight=nothing, init_bias=nothing, use_bias=True()) ``` Create a traditional fully connected layer, whose forward pass is given by: `y = activation.(weight * x .+ bias)` **Arguments** * `in_dims`: number of input dimensions * `out_dims`: number of output dimensions * `activation`: activation function **Keyword Arguments** * `init_weight`: initializer for the weight matrix (`weight = init_weight(rng, out_dims, in_dims)`). If `nothing`, then we use [`kaiming_uniform`](/api/Building_Blocks/WeightInitializers#WeightInitializers.kaiming_uniform) with gain computed on the basis of the activation function (taken from Pytorch [`nn.init.calculate_gain`](https://pytorch.org/docs/stable/nn.init.html#torch.nn.init.calculate_gain)). * `init_bias`: initializer for the bias vector (ignored if `use_bias=false`). If `nothing`, then we use uniform distribution with bounds `-bound` and `bound` where `bound = inv(sqrt(in_dims))`. * `use_bias`: Trainable bias can be disabled entirely by setting this to `false` **Input** * `x` must be an AbstractArray with `size(x, 1) == in_dims` **Returns** * AbstractArray with dimensions `(out_dims, ...)` where `...` are the dimensions of `x` * Empty `NamedTuple()` **Parameters** * `weight`: Weight Matrix of size `(out_dims, in_dims)` * `bias`: Bias of size `(out_dims, 1)` (present if `use_bias=true`) source ```julia Scale(dims, activation=identity; init_weight=ones32, init_bias=zeros32, use_bias=True()) ``` Create a Sparsely Connected Layer with a very specific structure (only Diagonal Elements are non-zero). The forward pass is given by: `y = activation.(weight .* x .+ bias)` **Arguments** * `dims`: size of the learnable scale and bias parameters. * `activation`: activation function **Keyword Arguments** * `init_weight`: initializer for the weight matrix (`weight = init_weight(rng, out_dims, in_dims)`) * `init_bias`: initializer for the bias vector (ignored if `use_bias=false`) * `use_bias`: Trainable bias can be disabled entirely by setting this to `false` **Input** * `x` must be an Array of size `(dims..., B)` or `(dims...[0], ..., dims[k])` for `k ≤ size(dims)` **Returns** * Array of size `(dims..., B)` or `(dims...[0], ..., dims[k])` for `k ≤ size(dims)` * Empty `NamedTuple()` **Parameters** * `weight`: Weight Array of size `(dims...)` * `bias`: Bias of size `(dims...)` source ## Attention Layers {#Attention-Layers} ```julia MultiHeadAttention(dims; nheads=1, dense_kwargs=(; use_bias=False()), attention_dropout_probability=0.0f0, is_causal::Union{Bool,Nothing}=nothing) ``` The multi-head dot-product attention layer used in Transformer architectures \[[1](/references#vaswani2017attention)]. **Arguments** * `dims`: The embedding dimensions of inputs, intermediate tensors and outputs. In the most general case, it is given as * a) `(q_in_dim, k_in_dim, v_in_dim) => (qk_dim, v_dim) => out_dim`. Can take also simpler forms as * b) `dims::Int`; * c) `in_dim::Int => (qk_dim, v_dim) => out_dim`; * d) `in_dim::Int => qkv_dim => out_dim`. **Keyword Arguments** * `nheads`: number of heads. * `attention_dropout_probability`: dropout probability for the attention scores. * `dense_kwargs`: keyword arguments for the Dense layers. Default `use_bias=false`. * `is_causal`: whether the attention is causal. If this is provided, the attention mask will be automatically created (passing in the `mask` argument is not allowed and will throw an error). **Forward Pass Signature(s)** ```julia (m::MultiHeadAttention)(qkv, ps, st::NamedTuple) (m::MultiHeadAttention)((q, kv), ps, st::NamedTuple) (m::MultiHeadAttention)((q, k, v, [mask = nothing]), ps, st::NamedTuple) ``` **Inputs** * `qkv`: a single input tensor for query, key and value. This corresponds to self-attention. * `(q, kv)`: a tuple of two input tensors for query and key-value. * `(q, k, v)`: a tuple of three input tensors for query, key and value. * `mask`: an optional mask to apply to the attention scores. This must be broadcastable to the shape of the attention scores `(kv_len, q_len, nheads, batch_size)`. The query tensor `q` is expected to have shape `(q_in_dim, q_len, batch_size)`, the key test `k` is expected to have shape `(k_in_dim, kv_len, batch_size)`, the value tensor `v` is expected to have shape `(v_in_dim, kv_len, batch_size)`. **Returns** * A tuple of two elements. The first element is the output tensor of shape `(out_dim, q_len, batch_size)` and the second element is the attention scores of shape `(q_len, kv_len, nheads, batch_size)`. * A NamedTuple of the states of the layer. **Extended Help** **Examples** ```julia julia> m = MultiHeadAttention(64; nheads=8); julia> ps, st = Lux.setup(Random.default_rng(), m); julia> q = randn(Float32, 64, 10, 32); julia> k = randn(Float32, 64, 20, 32); julia> v = randn(Float32, 64, 20, 32); julia> (y, α), st_new = m((q, k, v), ps, st); julia> size(y) (64, 10, 32) julia> size(α) (20, 10, 8, 32) julia> (y, α), st_new = m(q, ps, st); # self-attention julia> size(y) (64, 10, 32) julia> size(α) (10, 10, 8, 32) ``` source ## Embedding Layers {#Embedding-Layers} ```julia Embedding(in_dims => out_dims; init_weight=rand32) ``` A lookup table that stores embeddings of dimension `out_dims` for a vocabulary of size `in_dims`. When the vocabulary is multi-dimensional, the input is expected to be a tuple of Cartesian indices. This layer is often used to store word embeddings and retrieve them using indices. **Arguments** * `in_dims`: number(s) of input dimensions * `out_dims`: number of output dimensions **Keyword Arguments** * `init_weight`: initializer for the weight matrix (`weight = init_weight(rng, out_dims, in_dims...)`) **Input** * Integer OR * Abstract Vector of Integers OR * Abstract Array of Integers OR * Tuple of Integers OR * Tuple of Abstract Vectors of Integers OR * Tuple of Abstract Arrays of Integers **Returns** * Returns the embedding corresponding to each index in the input. For an N dimensional input, an N + 1 dimensional output is returned. * Empty `NamedTuple()` ::: warning Gradients with Tracker.jl Tracker.jl produces incorrect gradients for this layer if indices in the input are repeated. Don't use this layer with Tracker.jl if you need to compute gradients. ::: source ```julia RotaryPositionalEmbedding( dim::IntegerType; max_sequence_length::IntegerType=4096, base::IntegerType=10000, low_memory_variant::Bool=true, ) ``` Rotary Positional Embedding. For details see [Su *et al.* \[2\]](/references#su2024roformer). The traditional implementation rotates consecutive pairs of elements in the feature dimension while the default implementation rotates pairs with stride half the feature dimensions for efficiency. **Arguments** * `dim`: The feature dimensions to be rotated. If the input feature is larger than dims then the rest is left unchanged. **Keyword Arguments** * `base`: The base used to compute angular frequency for each dimension in the positional encodings. Default: `10000`. * `max_sequence_length`: The maximum sequence length. Default: `4096`. * `low_memory_variant`: If `true` then cos and sin cache have leading dimension of `dim ÷ 2`. If `false` then cos and sin cache have leading dimension of `dim`. Default: `true`. **Input** * 4D `AbstractArray` such that * `size(x, 1) == dim` * `size(x, 3) ≤ max_sequence_length` **Returns** * 4D `AbstractArray` of the same size as the input. **States** * NamedTuple containing `cos_cache` and `sin_cache`. source ```julia SinusoidalPositionalEmbedding( dims::IntegerType; min_freq=0.0001f0, max_freq=1.0f0, scale=nothing, full_turns::Bool=false ) ``` Sinusoidal Positional Embedding. For details see [Vaswani *et al.* \[1\]](/references#vaswani2017attention). **Arguments** * `dims`: The dimensionality of the resulting positional embeddings. **Keyword Arguments** * `min_freq`: The minimum frequency expected. Default: `0.0001f0`. * `max_freq`: The maximum frequency expected. Default: `1.0f0`. * `scale`: A multiplicative scale for the embeddings. Default: $\sqrt{2/dims}$. * `full_turns`: If `true` multiply the frequencies with $2\pi$. Default: `false`. **Input** * AbstractArray **Returns** * If the input array is of size `(insz...,)` then the output is of size `(dims, insz...)`. **States** * NamedTuple containing `sigmas`. source ### Functional API {#Functional-API} ```julia apply_rotary_embedding(x::AbstractArray{T,4}, cos_cache::AbstractMatrix, sin_cache::AbstractMatrix; head_dim::Integer, seq_dim::Integer) apply_rotary_embedding(x::AbstractArray{T,4}, input_positions::AbstractVector{<:Integer}, cos_cache::AbstractMatrix, sin_cache::AbstractMatrix; head_dim::Integer, seq_dim::Integer) ``` Apply rotary embedding to the input `x` using the `cos_cache` and `sin_cache` parameters. If `input_positions` is provided, then we extract the cosine and sine cache for the corresponding positions in the sequence. Otherwise, we use the entire cache upto sequence length of `x`. **Arguments** * `x`: 4D `AbstractArray`. * `cos_cache`: Cache of cosine values. Generated using [`compute_rotary_embedding_params`](/api/Lux/layers#Lux.compute_rotary_embedding_params). * `sin_cache`: Cache of sine values. Generated using [`compute_rotary_embedding_params`](/api/Lux/layers#Lux.compute_rotary_embedding_params). * `seq_dim`: Dimension of the sequence. Must be between 1 and 4. * `input_positions`: Positions in the sequence to extract the cosine and sine cache for. If not provided, then we use the entire cache upto sequence length of `x`. **Returns** * Output of the rotary embedding. source ```julia compute_rotary_embedding_params(head_dim::Integer, max_sequence_length::Integer; base::Number, dtype::Type{T}=Float32, low_memory_variant::Bool=true) ``` Computes the cosine and sine cache for rotary positional embeddings. **Arguments** * `head_dim`: The feature dimensions to be rotated. * `max_sequence_length`: The maximum sequence length. Default: `4096`. **Keyword Arguments** * `base`: The base used to compute angular frequency for each dimension in the positional encodings. Default: `10000`. * `dtype`: The data type of the cache. Default: `Float32`. * `low_memory_variant`: If `true` then cos and sin cache have leading dimension of `head_dim ÷ 2`. If `false` then cos and sin cache have leading dimension of `head_dim`. Default: `true`. source ## Misc. Helper Layers {#Misc.-Helper-Layers} ```julia FlattenLayer(; N = nothing) ``` Flattens the passed array into a matrix. **Keyword Arguments** * `N`: Flatten the first `N` dimensions of the input array. If `nothing`, then all dimensions (except the last) are flattened. Note that the batch dimension is never flattened. **Inputs** * `x`: AbstractArray **Returns** * AbstractMatrix of size `(:, size(x, ndims(x)))` if `N` is `nothing` else the first `N` dimensions of the input array are flattened. * Empty `NamedTuple()` **Example** ```julia julia> model = FlattenLayer() FlattenLayer{Nothing}(nothing) julia> rng = Random.default_rng(); Random.seed!(rng, 0); ps, st = Lux.setup(rng, model); x = randn(rng, Float32, (2, 2, 2, 2)); julia> y, st_new = model(x, ps, st); size(y) (8, 2) ``` source ```julia Maxout(layers...) Maxout(; layers...) Maxout(f::Function, n_alts::Int) ``` This contains a number of internal layers, each of which receives the same input. Its output is the elementwise maximum of the the internal layers' outputs. Maxout over linear dense layers satisfies the universal approximation theorem \[[3](/references#goodfellow2013maxout)]. See also [`Parallel`](/api/Lux/layers#Lux.Parallel) to reduce with other operators. **Arguments** * Layers can be specified in three formats: * A list of `N` Lux layers * Specified as `N` keyword arguments. * A no argument function `f` and an integer `n_alts` which specifies the number of layers. **Extended Help** **Inputs** * `x`: Input that is passed to each of the layers **Returns** * Output is computed by taking elementwise `max` of the outputs of the individual layers. * Updated state of the `layers` **Parameters** * Parameters of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) **States** * States of each `layer` wrapped in a NamedTuple with `fields = layer_1, layer_2, ..., layer_N` (naming changes if using the kwargs API) source ```julia NoOpLayer() ``` As the name suggests does nothing but allows pretty printing of layers. Whatever input is passed is returned. **Example** ```julia julia> model = NoOpLayer() NoOpLayer() julia> rng = Random.default_rng(); Random.seed!(rng, 0); ps, st = Lux.setup(rng, model); x = 1 1 julia> y, st_new = model(x, ps, st) (1, NamedTuple()) ``` source ```julia ReshapeLayer(dims) ``` Reshapes the passed array to have a size of `(dims..., :)` **Arguments** * `dims`: The new dimensions of the array (excluding the last dimension). **Inputs** * `x`: AbstractArray of any shape which can be reshaped in `(dims..., size(x, ndims(x)))` **Returns** * AbstractArray of size `(dims..., size(x, ndims(x)))` * Empty `NamedTuple()` **Example** ```julia julia> model = ReshapeLayer((2, 2)) ReshapeLayer(output_dims = (2, 2, :)) julia> rng = Random.default_rng(); Random.seed!(rng, 0); ps, st = Lux.setup(rng, model); x = randn(rng, Float32, (4, 1, 3)); julia> y, st_new = model(x, ps, st); size(y) (2, 2, 3) ``` source ```julia SelectDim(dim, i) ``` Return a view of all the data of the input `x` where the index for dimension `dim` equals `i`. Equivalent to `view(x,:,:,...,i,:,:,...)` where `i` is any valid index for index slot `d` (e.g. an integer or a unit range). Note that it may be inefficient to use non-contiguous views. **Arguments** * `dim`: Dimension for indexing * `i`: Index or indices for dimension `dim` **Inputs** * `x`: AbstractArray that can be indexed with `view(x,:,:,...,i,:,:,...)` **Returns** * `view(x,:,:,...,i,:,:,...)` where `i` is in position `d` * Empty `NamedTuple()` source ```julia WrappedFunction(f) ``` Wraps a stateless and parameter less function. Might be used when a function is added to `Chain`. For example, `Chain(x -> relu.(x))` would not work and the right thing to do would be `Chain((x, ps, st) -> (relu.(x), st))`. An easier thing to do would be `Chain(WrappedFunction(Base.Fix1(broadcast, relu)))` **Arguments** * `f`: Some function. **Inputs** * `x`: will be directly passed to `f` **Returns** * Output of `f(x)` * Empty `NamedTuple()` source ```julia ReverseSequence(dim = nothing) ``` Reverse the specified dimension `dims` of the passed array **Arguments** * `dim`: Dimension that need to be reversed. If `nothing`, for AbstractVector{T} it reverses itself (dimension 1), for other arrays, reverse the dimension `ndims(x) - 1`. **Inputs** * `x`: AbstractArray. **Returns** * AbstractArray with the same dimensions as the input * Empty `NamedTuple()` **Example** ```julia julia> model = ReverseSequence() ReverseSequence{Nothing}(nothing) julia> rng = Random.default_rng(); Random.seed!(rng, 0); ps, st = Lux.setup(rng, model); x = [1.0, 2.0, 3.0]; julia> y, st_new = model(x, ps, st) ([3.0, 2.0, 1.0], NamedTuple()) ``` source ## Normalization Layers {#Normalization-Layers} ```julia BatchNorm(chs::Integer, activation=identity; init_bias=zeros32, init_scale=ones32, affine=True(), track_stats=True(), epsilon=1f-5, momentum=0.1f0, use_decomposed_implementation=False()) ``` [Batch Normalization](https://arxiv.org/abs/1502.03167) layer. `BatchNorm` computes the mean and variance for each $D\_1 × ... × D\_{N-2} × 1 × D\_N$ input slice and normalises the input accordingly. **Arguments** * `chs`: Size of the channel dimension in your data. Given an array with `N` dimensions, call the `N-1`th the channel dimension. For a batch of feature vectors this is just the data dimension, for `WHCN` images it's the usual channel dimension. * `activation`: After normalization, elementwise activation `activation` is applied. **Keyword Arguments** * If `track_stats=true`, accumulates mean and variance statistics in training phase that will be used to renormalize the input in test phase. * `epsilon`: a value added to the denominator for numerical stability * `momentum`: the value used for the `running_mean` and `running_var` computation * If `affine=true`, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters. * `init_bias`: Controls how the `bias` is initialized * `init_scale`: Controls how the `scale` is initialized **Extended Help** **Internal Keyword Arguments** * `use_decomposed_implementation`: Several backends like CUDA, Reactant dispatch to a specialized vendored implementation for batchnorm. Setting this to true, makes Lux emit a variant of batchnorm that is computed without using any specialized kernels. This is meant to be used for correctness testing and benchmarking purposes. **Inputs** * `x`: Array where `size(x, N - 1) = chs` **Returns** * `y`: Normalized Array * Update model state **Parameters** * `affine=true` * `bias`: Bias of shape `(chs,)` * `scale`: Scale of shape `(chs,)` * `affine=false` - Empty `NamedTuple()` **States** * Statistics if `track_stats=true` * `running_mean`: Running mean of shape `(chs,)` * `running_var`: Running variance of shape `(chs,)` * Statistics if `track_stats=false` * `running_mean`: nothing * `running_var`: nothing * `training`: Used to check if training/inference mode Use `Lux.testmode` during inference. **Example** ```julia julia> Chain(Dense(784 => 64), BatchNorm(64, relu), Dense(64 => 10), BatchNorm(10)) Chain( layer_1 = Dense(784 => 64), # 50_240 parameters layer_2 = BatchNorm(64, relu, affine=true, track_stats=true), # 128 parameters, plus 129 non-trainable layer_3 = Dense(64 => 10), # 650 parameters layer_4 = BatchNorm(10, affine=true, track_stats=true), # 20 parameters, plus 21 non-trainable ) # Total: 51_038 parameters, # plus 150 states. ``` ::: warning Warning Passing a batch size of 1, during training will result in an error. ::: See also [`BatchNorm`](/api/Lux/layers#Lux.BatchNorm), [`InstanceNorm`](/api/Lux/layers#Lux.InstanceNorm), [`LayerNorm`](/api/Lux/layers#Lux.LayerNorm), [`WeightNorm`](/api/Lux/layers#Lux.WeightNorm) source ```julia GroupNorm(chs::Integer, groups::Integer, activation=identity; init_bias=zeros32, init_scale=ones32, affine=true, epsilon=1f-5) ``` [Group Normalization](https://arxiv.org/abs/1803.08494) layer. **Arguments** * `chs`: Size of the channel dimension in your data. Given an array with `N` dimensions, call the `N-1`th the channel dimension. For a batch of feature vectors this is just the data dimension, for `WHCN` images it's the usual channel dimension. * `groups` is the number of groups along which the statistics are computed. The number of channels must be an integer multiple of the number of groups. * `activation`: After normalization, elementwise activation `activation` is applied. **Keyword Arguments** * `epsilon`: a value added to the denominator for numerical stability * If `affine=true`, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters. * `init_bias`: Controls how the `bias` is initialized * `init_scale`: Controls how the `scale` is initialized **Extended Help** **Inputs** * `x`: Array where `size(x, N - 1) = chs` and `ndims(x) > 2` **Returns** * `y`: Normalized Array * Update model state **Parameters** * `affine=true` * `bias`: Bias of shape `(chs,)` * `scale`: Scale of shape `(chs,)` * `affine=false` - Empty `NamedTuple()` **States** * `training`: Used to check if training/inference mode Use `Lux.testmode` during inference. **Example** ```julia julia> Chain(Dense(784 => 64), GroupNorm(64, 4, relu), Dense(64 => 10), GroupNorm(10, 5)) Chain( layer_1 = Dense(784 => 64), # 50_240 parameters layer_2 = GroupNorm(64, 4, relu, affine=true), # 128 parameters layer_3 = Dense(64 => 10), # 650 parameters layer_4 = GroupNorm(10, 5, affine=true), # 20 parameters ) # Total: 51_038 parameters, # plus 0 states. ``` See also [`GroupNorm`](/api/Lux/layers#Lux.GroupNorm), [`InstanceNorm`](/api/Lux/layers#Lux.InstanceNorm), [`LayerNorm`](/api/Lux/layers#Lux.LayerNorm), [`WeightNorm`](/api/Lux/layers#Lux.WeightNorm) source ```julia InstanceNorm(chs::Integer, activation=identity; init_bias=zeros32, init_scale=ones32, affine=False(), track_stats=False(), epsilon=1f-5, momentum=0.1f0) ``` Instance Normalization. For details see [Ulyanov *et al.* \[4\]](/references#ulyanov2016instance). Instance Normalization computes the mean and variance for each $D\_1 \times ... \times D\_{N - 2} \times 1 \times 1$\` input slice and normalises the input accordingly. **Arguments** * `chs`: Size of the channel dimension in your data. Given an array with `N` dimensions, call the `N-1`th the channel dimension. For a batch of feature vectors this is just the data dimension, for `WHCN` images it's the usual channel dimension. * `activation`: After normalization, elementwise activation `activation` is applied. **Keyword Arguments** * If `track_stats=true`, accumulates mean and variance statistics in training phase that will be used to renormalize the input in test phase. * `epsilon`: a value added to the denominator for numerical stability * `momentum`: the value used for the `running_mean` and `running_var` computation * If `affine=true`, it also applies a shift and a rescale to the input through to learnable per-channel bias and scale parameters. * `init_bias`: Controls how the `bias` is initialized * `init_scale`: Controls how the `scale` is initialized **Extended Help** **Inputs** * `x`: Array where `size(x, N - 1) = chs` and `ndims(x) > 2` **Returns** * `y`: Normalized Array * Update model state **Parameters** * `affine=true` * `bias`: Bias of shape `(chs,)` * `scale`: Scale of shape `(chs,)` * `affine=false` - Empty `NamedTuple()` **States** * Statistics if `track_stats=true` * `running_mean`: Running mean of shape `(chs,)` * `running_var`: Running variance of shape `(chs,)` * Statistics if `track_stats=false` * `running_mean`: nothing * `running_var`: nothing * `training`: Used to check if training/inference mode Use `Lux.testmode` during inference. **Example** ```julia julia> Chain(Dense(784 => 64), InstanceNorm(64, relu; affine=true), Dense(64 => 10), InstanceNorm(10, relu; affine=true)) Chain( layer_1 = Dense(784 => 64), # 50_240 parameters layer_2 = InstanceNorm(64, relu, affine=true, track_stats=false), # 128 parameters, plus 1 non-trainable layer_3 = Dense(64 => 10), # 650 parameters layer_4 = InstanceNorm(10, relu, affine=true, track_stats=false), # 20 parameters, plus 1 non-trainable ) # Total: 51_038 parameters, # plus 2 states. ``` See also [`BatchNorm`](/api/Lux/layers#Lux.BatchNorm), [`GroupNorm`](/api/Lux/layers#Lux.GroupNorm), [`LayerNorm`](/api/Lux/layers#Lux.LayerNorm), [`WeightNorm`](/api/Lux/layers#Lux.WeightNorm) source ```julia LayerNorm(shape::NTuple{N, Int}, activation=identity; epsilon=1f-5, dims=Colon(), affine=true, init_bias=zeros32, init_scale=ones32) ``` Computes mean and standard deviation over the whole input array, and uses these to normalize the whole array. Optionally applies an elementwise affine transformation afterwards. Given an input array $x$, this layer computes $$y = \frac{x - \mathbb{E}\[x]}{\sqrt{Var\[x] + \epsilon}} \* \gamma + \beta$$ where $\gamma$ & $\beta$ are trainable parameters if `affine=true`. **Arguments** * `shape`: Broadcastable shape of input array excluding the batch dimension. * `activation`: After normalization, elementwise activation `activation` is applied. **Keyword Arguments** * `epsilon`: a value added to the denominator for numerical stability. * `dims`: Dimensions to normalize the array over. * If `affine=true`, it also applies a shift and a rescale to the input through to learnable per-element bias and scale parameters. * `init_bias`: Controls how the `bias` is initialized * `init_scale`: Controls how the `scale` is initialized **Extended Help** **Inputs** * `x`: AbstractArray **Returns** * `y`: Normalized Array * Empty NamedTuple() **Parameters** * `affine=false`: Empty `NamedTuple()` * `affine=true` * `bias`: Bias of shape `(shape..., 1)` * `scale`: Scale of shape `(shape..., 1)` source ```julia WeightNorm(layer::AbstractLuxLayer, which_params::NTuple{N, Symbol}, dims::Union{Tuple, Nothing}=nothing) ``` Applies [weight normalization](https://arxiv.org/abs/1602.07868) to a parameter in the given layer. $$w = g\frac{v}{|v|}$$ Weight normalization is a reparameterization that decouples the magnitude of a weight tensor from its direction. This updates the parameters in `which_params` (e.g. `weight`) using two parameters: one specifying the magnitude (e.g. `weight_g`) and one specifying the direction (e.g. `weight_v`). **Arguments** * `layer` whose parameters are being reparameterized * `which_params`: parameter names for the parameters being reparameterized * By default, a norm over the entire array is computed. Pass `dims` to modify the dimension. **Inputs** * `x`: Should be of valid type for input to `layer` **Returns** * Output from `layer` * Updated model state of `layer` **Parameters** * `normalized`: Parameters of `layer` that are being normalized * `unnormalized`: Parameters of `layer` that are not being normalized **States** * Same as that of `layer` source ```julia RMSNorm(normalized_shape::Dims; epsilon=1.0f-5, affine=true) RMSNorm(dim::Integer...; kwargs...) ``` Root Mean Square Normalization layer. It normalizes the input by computing the root mean square (RMS) of the first `N` dimensions of the input where `N` is the length of `normalized_shape`. $$\begin{align} y\_i &= \frac{x\_i}{\sqrt{\epsilon + \frac{1}{D}\Sigma\_{j=1}^D x\_j^2}} \* \gamma\_i \end{align}$$ **Arguments** * `normalized_shape`: The input shape from which the RMS normalization factor is computed. The input is expected to have a shape that can be broadcast with `normalized_shape`. **Keyword Arguments** * `epsilon`: A small value for numerical stability. * `affine`: If `true`, learns a scale parameter. **Extended Help** **Inputs** * `x`: Array of size `(normalized_shape..., *, *..., *)` **Returns** * `y`: Normalized Array of same shape as `x` * Empty `NamedTuple()` source ## Upsampling {#Upsampling} ```julia PixelShuffle(r::Int) ``` Pixel shuffling layer with upscale factor `r`. Usually used for generating higher resolution images while upscaling them. See `NNlib.pixel_shuffle` for more details. **Arguments** * `r`: Upscale factor **Inputs** * `x`: For 4D-arrays representing N images, the operation converts input `size(x) == (W, H, r² x C, N)` to output of size `(r x W, r x H, C, N)`. For D-dimensional data, it expects `ndims(x) == D + 2` with channel and batch dimensions, and divides the number of channels by `rᴰ`. **Returns** * Output of size `(r x W, r x H, C, N)` for 4D-arrays, and `(r x W, r x H, ..., C, N)` for D-dimensional data, where `D = ndims(x) - 2` source ```julia Upsample(mode = :nearest; [scale, size, align_corners=false]) Upsample(scale, mode = :nearest) ``` Upsampling Layer. **Layer Construction** **Option 1** * `mode`: Set to `:nearest`, `:linear`, `:bilinear` or `:trilinear` Exactly one of two keywords must be specified: * If `scale` is a number, this applies to all but the last two dimensions (channel and batch) of the input. It may also be a tuple, to control dimensions individually. * Alternatively, keyword `size` accepts a tuple, to directly specify the leading dimensions of the output. **Option 2** * If `scale` is a number, this applies to all but the last two dimensions (channel and batch) of the input. It may also be a tuple, to control dimensions individually. * `mode`: Set to `:nearest`, `:bilinear` or `:trilinear` Currently supported upsampling `mode`s and corresponding NNlib's methods are: * `:nearest` -> `NNlib.upsample_nearest` * `:bilinear` -> `NNlib.upsample_bilinear` * `:trilinear` -> `NNlib.upsample_trilinear` **Extended Help** **Other Keyword Arguments** * `align_corners`: If `true`, the corner pixels of the input and output tensors are aligned, and thus preserving the values at those pixels. This only has effect when mode is one of `:bilinear` or `:trilinear`. **Inputs** * `x`: For the input dimensions look into the documentation for the corresponding `NNlib` function * As a rule of thumb, `:nearest` should work with arrays of arbitrary dimensions * `:bilinear` works with 4D Arrays * `:trilinear` works with 5D Arrays **Returns** * Upsampled Input of size `size` or of size `(I_1 x scale[1], ..., I_N x scale[N], C, N)` * Empty `NamedTuple()` source --- --- url: /dev/api/Lux/autodiff.md --- # Automatic Differentiation Helpers {#autodiff-lux-helpers} ## JVP & VJP Wrappers {#JVP-and-VJP-Wrappers} ```julia jacobian_vector_product(f, backend::AbstractADType, x, u) ``` Compute the Jacobian-Vector Product $\left(\frac{\partial f}{\partial x}\right) u$. This is a wrapper around AD backends but allows us to compute gradients of jacobian-vector products efficiently using mixed-mode AD. **Backends & AD Packages** | Supported Backends | Packages Needed | Notes | |:------------------ |:--------------------------- |:------------------------------------------------- | | `AutoEnzyme` | `Enzyme.jl` / `Reactant.jl` | For nested AD support directly using `Enzyme.jl`. | | `AutoForwardDiff` | | | ::: warning Only for ChainRules-based AD like Zygote Gradient wrt `u` in the reverse pass is always dropped. ::: **Arguments** * `f`: The function to compute the jacobian of. * `backend`: The backend to use for computing the JVP. * `x`: The input to the function. * `u`: An object of the same structure as `x`. **Returns** * `v`: The Jacobian Vector Product. source ```julia vector_jacobian_product(f, backend::AbstractADType, x, u) ``` Compute the Vector-Jacobian Product $\left(\frac{\partial f}{\partial x}\right)^T u$. This is a wrapper around AD backends but allows us to compute gradients of vector-jacobian products efficiently using mixed-mode AD. **Backends & AD Packages** | Supported Backends | Packages Needed | Notes | |:------------------ |:--------------------------- |:------------------------------------------------- | | `AutoEnzyme` | `Enzyme.jl` / `Reactant.jl` | For nested AD support directly using `Enzyme.jl`. | | `AutoZygote` | `Zygote.jl` | | ::: warning Only for ChainRules-based AD like Zygote Gradient wrt `u` in the reverse pass is always dropped. ::: **Arguments** * `f`: The function to compute the jacobian of. * `backend`: The backend to use for computing the VJP. * `x`: The input to the function. * `u`: An object of the same structure as `f(x)`. **Returns** * `v`: The Vector Jacobian Product. source ## Batched AD {#Batched-AD} ```julia batched_jacobian(f, backend::AbstractADType, x::AbstractArray) ``` Computes the Jacobian of a function `f` with respect to a batch of inputs `x`. This expects the following properties for `y = f(x)`: 1. `ndims(y) ≥ 2` 2. `size(y, ndims(y)) == size(x, ndims(x))` **Backends & AD Packages** | Supported Backends | Packages Needed | |:------------------ |:--------------- | | `AutoEnzyme` | `Reactant.jl` | | `AutoForwardDiff` | | | `AutoZygote` | `Zygote.jl` | **Arguments** * `f`: The function to compute the jacobian of. * `backend`: The backend to use for computing the jacobian. * `x`: The input to the function. Must have `ndims(x) ≥ 2`. **Returns** * `J`: The Jacobian of `f` with respect to `x`. This will be a 3D Array. If the dimensions of `x` are `(N₁, N₂, ..., Nₙ, B)` and of `y` are `(M₁, M₂, ..., Mₘ, B)`, then `J` will be a `((M₁ × M₂ × ... × Mₘ), (N₁ × N₂ × ... × Nₙ), B)` Array. ::: danger Danger `f(x)` must not be inter-mixing the batch dimensions, else the result will be incorrect. For example, if `f` contains operations like batch normalization, then the result will be incorrect. ::: source ## Nested 2nd Order AD {#Nested-2nd-Order-AD} Consult the [manual page on Nested AD](/manual/nested_autodiff#nested_autodiff) for information on nested automatic differentiation. --- --- url: /dev/api/Lux/utilities.md --- # Utilities {#Utilities} ## Training API {#Training-API} Helper Functions making it easier to train `Lux.jl` models. Training is meant to be simple and provide extremely basic functionality. We provide basic building blocks which can be seamlessly composed to create complex training pipelines. ```julia TrainState ``` Training State containing: * `model`: `Lux` model. * `parameters`: Trainable Variables of the `model`. * `states`: Non-trainable Variables of the `model`. * `optimizer`: Optimizer from `Optimisers.jl`. * `optimizer_state`: Optimizer State. * `step`: Number of updates of the parameters made. Internal fields: * `cache`: Cached values. Implementations are free to use this for whatever they want. * `allocator_cache`: Used by GPUArrays compatible backends to cache memory allocations. * `objective_function`: Objective function might be cached. ::: warning Warning Constructing this object directly shouldn't be considered a stable API. Use the version with the Optimisers API. ::: source ```julia TrainState(model::Lux.AbstractLuxLayer, ps, st, optimizer::Optimisers.AbstractRule) ``` Constructor for [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState). **Arguments** * `ps`: Parameters of the model. * `st`: States of the model. * `model`: `Lux` model. * `optimizer`: Optimizer from `Optimisers.jl`. **Returns** [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState) object. source ```julia compute_gradients( ad::AbstractADType, objective_function::Function, data, ts::TrainState; sync::Bool=false, compile_options::Union{Missing,Reactant.CompileOptions}=missing ) ``` Compute the gradients of the objective function wrt parameters stored in `ts`. **Backends & AD Packages** | Supported Backends | Packages Needed | |:---------------------------- |:---------------- | | `AutoZygote` | `Zygote.jl` | | `AutoReverseDiff(; compile)` | `ReverseDiff.jl` | | `AutoTracker` | `Tracker.jl` | | `AutoEnzyme` | `Enzyme.jl` | | `AutoForwardDiff` | | | `AutoMooncake` | `Mooncake.jl` | **Arguments** * `ad`: Backend (from [ADTypes.jl](https://github.com/SciML/ADTypes.jl)) used to compute the gradients. * `objective_function`: Objective function. The function must take 4 inputs – model, parameters, states and data. The function must return 3 values – loss, updated\_state, and any computed statistics. * `data`: Data used to compute the gradients. * `ts`: Current Training State. See [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState). **Keyword Arguments** * `sync`: If `true`, then the compiled reactant function is compiled with `sync=true`. Typically reactant functions are asynchronous, which means if used with profiling or for timing, the timing will be inaccurate. Setting `sync=true` will ensure that the function will finish execution before this function returns. This is only used for Reactant Backend. * `compile_options`: Compile options for the reactant function. See `Reactant.CompileOptions` for more details. This is only used for Reactant Backend. **Return** A 4-Tuple containing: * `grads`: Computed Gradients. * `loss`: Loss from the objective function. * `stats`: Any computed statistics from the objective function. * `ts`: Updated Training State. **Known Limitations** * `AutoReverseDiff(; compile=true)` is not supported for Lux models with non-empty state `st`. Additionally the returned stats must be empty (`NamedTuple()`). We catch these issues in most cases and throw an error. * AutoForwardDiff only works with parameters that are AbstractArrays (e.g. ps=ComponentVector(ps)) ::: danger Aliased Gradients `grads` returned by this function might be aliased by the implementation of the gradient backend. For example, if you cache the `grads` from step `i`, the new gradients returned in step `i + 1` might be aliased by the old gradients. If you want to prevent this, simply use `copy(grads)` or `deepcopy(grads)` to make a copy of the gradients. ::: source ```julia apply_gradients(ts::TrainState, grads) ``` Update the parameters stored in `ts` using the gradients `grads`. **Arguments** * `ts`: [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState) object. * `grads`: Gradients of the loss function wrt `ts.params`. **Returns** Updated [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState) object. source ```julia apply_gradients!(ts::TrainState, grads) ``` Update the parameters stored in `ts` using the gradients `grads`. This is an inplace version of [`apply_gradients`](/api/Lux/utilities#Lux.Training.apply_gradients). **Arguments** * `ts`: [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState) object. * `grads`: Gradients of the loss function wrt `ts.params`. **Returns** Updated [`TrainState`](/api/Lux/utilities#Lux.Training.TrainState) object. source ```julia single_train_step( backend, obj_fn::F, data, ts::TrainState; return_gradients=True(), sync::Bool=false, compile_options::Union{Nothing,Reactant.CompileOptions}=missing, ) ``` Perform a single training step. Computes the gradients using [`compute_gradients`](/api/Lux/utilities#Lux.Training.compute_gradients) and updates the parameters using [`apply_gradients`](/api/Lux/utilities#Lux.Training.apply_gradients). All backends supported via [`compute_gradients`](/api/Lux/utilities#Lux.Training.compute_gradients) are supported here. In most cases you should use [`single_train_step!`](/api/Lux/utilities#Lux.Training.single_train_step!) instead of this function. **Keyword Arguments** * `return_gradients`: If `True()`, the gradients are returned. If `False()`, the returned gradients are `nothing`. Defaults to `True()`. This is only used for Reactant Backend. * `sync`: If `true`, then the compiled reactant function is compiled with `sync=true`. Typically reactant functions are asynchronous, which means if used with profiling or for timing, the timing will be inaccurate. Setting `sync=true` will ensure that the function will finish execution before this function returns. This is only used for Reactant Backend. * `compile_options`: Compile options for the reactant function. See `Reactant.CompileOptions` for more details. This is only used for Reactant Backend. **Return** Returned values are the same as [`single_train_step!`](/api/Lux/utilities#Lux.Training.single_train_step!). source ```julia single_train_step!( backend, obj_fn::F, data, ts::TrainState; return_gradients=True(), sync::Bool=false, compile_options::Union{Nothing,Reactant.CompileOptions}=missing, ) ``` Perform a single training step. Computes the gradients using [`compute_gradients`](/api/Lux/utilities#Lux.Training.compute_gradients) and updates the parameters using [`apply_gradients!`](/api/Lux/utilities#Lux.Training.apply_gradients!). All backends supported via [`compute_gradients`](/api/Lux/utilities#Lux.Training.compute_gradients) are supported here. **Keyword Arguments** * `return_gradients`: If `True()`, the gradients are returned. If `False()`, the returned gradients are `nothing`. Defaults to `True()`. This is only used for Reactant Backend. * `sync`: If `true`, then the compiled reactant function is compiled with `sync=true`. Typically reactant functions are asynchronous, which means if used with profiling or for timing, the timing will be inaccurate. Setting `sync=true` will ensure that the function will finish execution before this function returns. This is only used for Reactant Backend. * `compile_options`: Compile options for the reactant function. See `Reactant.CompileOptions` for more details. This is only used for Reactant Backend. **Return** Returned values are the same as [`compute_gradients`](/api/Lux/utilities#Lux.Training.compute_gradients). Note that despite the `!`, only the parameters in `ts` are updated inplace. Users should be using the returned `ts` object for further training steps, else there is no caching and performance will be suboptimal (and absolutely terrible for backends like `AutoReactant`). source ## Loss Functions {#Loss-Functions} Loss Functions Objects take 2 forms of inputs: 1. `ŷ` and `y` where `ŷ` is the predicted output and `y` is the target output. 2. `model`, `ps`, `st`, `(x, y)` where `model` is the model, `ps` are the parameters, `st` are the states and `(x, y)` are the input and target pair. Then it returns the loss, updated states, and an empty named tuple. This makes them compatible with the [Training API](/api/Lux/utilities#Training-API). ::: warning Warning When using ChainRules.jl compatible AD (like Zygote), we only compute the gradients wrt the inputs and drop any gradients wrt the targets. ::: ```julia GenericLossFunction(loss_fn; agg = mean) ``` Takes any function `loss_fn` that maps 2 number inputs to a single number output. Additionally, array inputs are efficiently broadcasted and aggregated using `agg`. ```julia julia> mseloss = GenericLossFunction((ŷ, y) -> abs2(ŷ - y)); julia> y_model = [1.1, 1.9, 3.1]; julia> mseloss(y_model, 1:3) ≈ 0.01 true ``` **Special Note** This function takes any of the [`LossFunctions.jl`](https://juliaml.github.io/LossFunctions.jl/stable/) public functions into the Lux Losses API with efficient aggregation. source ```julia BinaryCrossEntropyLoss(; agg = mean, epsilon = nothing, label_smoothing::Union{Nothing, Real}=nothing, logits::Union{Bool, Val}=Val(false)) ``` Binary Cross Entropy Loss with optional label smoothing and fused logit computation. Returns the binary cross entropy loss computed as: * If `logits` is either `false` or `Val(false)`: $$\text{agg}\left(-\tilde{y} \* \log\left(\hat{y} + \epsilon\right) - (1 - \tilde{y}) \* \log\left(1 - \hat{y} + \epsilon\right)\right)$$ * If `logits` is `true` or `Val(true)`: $$\text{agg}\left((1 - \tilde{y}) \* \hat{y} - \log\sigma(\hat{y})\right)$$ The value of $\tilde{y}$ is computed using label smoothing. If `label_smoothing` is `nothing`, then no label smoothing is applied. If `label_smoothing` is a real number $\in \[0, 1]$, then the value of $\tilde{y}$ is: $$\tilde{y} = (1 - \alpha) \* y + \alpha \* 0.5$$ where $\alpha$ is the value of `label_smoothing`. **Extended Help** **Example** ```julia julia> bce = BinaryCrossEntropyLoss(); julia> y_bin = Bool[1, 0, 1]; julia> y_model = Float32[2, -1, pi] 3-element Vector{Float32}: 2.0 -1.0 3.1415927 julia> logitbce = BinaryCrossEntropyLoss(; logits=Val(true)); julia> logitbce(y_model, y_bin) ≈ 0.160832f0 true julia> bce(sigmoid.(y_model), y_bin) ≈ 0.16083185f0 true julia> bce_ls = BinaryCrossEntropyLoss(label_smoothing=0.1); julia> bce_ls(sigmoid.(y_model), y_bin) > bce(sigmoid.(y_model), y_bin) true julia> logitbce_ls = BinaryCrossEntropyLoss(label_smoothing=0.1, logits=Val(true)); julia> logitbce_ls(y_model, y_bin) > logitbce(y_model, y_bin) true ``` source ```julia BinaryFocalLoss(; gamma = 2, agg = mean, epsilon = nothing) ``` Return the binary focal loss \[[5](/references#lin2017focal)]. The model input, $\hat{y}$, is expected to be normalized (i.e. softmax output). For $\gamma = 0$ this is equivalent to [`BinaryCrossEntropyLoss`](/api/Lux/utilities#Lux.BinaryCrossEntropyLoss). **Example** ```julia julia> y = [0 1 0 1 0 1]; julia> ŷ = [0.268941 0.5 0.268941 0.731059 0.5 0.731059]; julia> BinaryFocalLoss()(ŷ, y) ≈ 0.0728675615927385 true julia> BinaryFocalLoss(gamma=0)(ŷ, y) ≈ BinaryCrossEntropyLoss()(ŷ, y) true ``` source ```julia CrossEntropyLoss(; agg=mean, epsilon=nothing, dims=1, logits::Union{Bool, Val}=Val(false), label_smoothing::Union{Nothing, Real}=nothing ) ``` Return the cross entropy loss which is used in multi-class classification tasks. The input, $\hat{y}$, is expected to be normalized (i.e. `softmax` output) if `logits` is `false` or `Val(false)`. The loss is calculated as: $$\text{agg}\left(-\sum \tilde{y} \log(\hat{y} + \epsilon)\right)$$ where $\epsilon$ is added for numerical stability. The value of $\tilde{y}$ is computed using label smoothing. If `label_smoothing` is `nothing`, then no label smoothing is applied. If `label_smoothing` is a real number $\in \[0, 1]$, then the value of $\tilde{y}$ is calculated as: $$\tilde{y} = (1 - \alpha) \* y + \alpha \* \text{size along dim}$$ where $\alpha$ is the value of `label_smoothing`. **Extended Help** **Example** ```julia julia> y = [1 0 0 0 1 0 1 0 1 0 0 0 1 0 0] 3×5 Matrix{Int64}: 1 0 0 0 1 0 1 0 1 0 0 0 1 0 0 julia> y_model = softmax(reshape(-7:7, 3, 5) .* 1f0) 3×5 Matrix{Float32}: 0.0900306 0.0900306 0.0900306 0.0900306 0.0900306 0.244728 0.244728 0.244728 0.244728 0.244728 0.665241 0.665241 0.665241 0.665241 0.665241 julia> CrossEntropyLoss()(y_model, y) ≈ 1.6076053f0 true julia> 5 * 1.6076053f0 ≈ CrossEntropyLoss(; agg=sum)(y_model, y) true julia> CrossEntropyLoss(label_smoothing=0.15)(y_model, y) ≈ 1.5776052f0 true ``` source ```julia DiceCoeffLoss(; smooth = true, agg = mean) ``` Return the Dice Coefficient loss \[[6](/references#milletari2016v)] which is used in segmentation tasks. The dice coefficient is similar to the F1\_score. Loss calculated as: $$\text{agg}\left(1 - \frac{2 \sum y \hat{y} + \alpha}{\sum y^2 + \sum \hat{y}^2 + \alpha}\right)$$ where $\alpha$ is the smoothing factor (`smooth`). **Example** ```julia julia> y_pred = [1.1, 2.1, 3.1]; julia> DiceCoeffLoss()(y_pred, 1:3) ≈ 0.000992391663909964 true julia> 1 - DiceCoeffLoss()(y_pred, 1:3) ≈ 0.99900760833609 true julia> DiceCoeffLoss()(reshape(y_pred, 3, 1), reshape(1:3, 3, 1)) ≈ 0.000992391663909964 true ``` source ```julia FocalLoss(; gamma = 2, dims = 1, agg = mean, epsilon = nothing) ``` Return the focal loss \[[5](/references#lin2017focal)] which can be used in classification tasks with highly imbalanced classes. It down-weights well-classified examples and focuses on hard examples. The input, $\hat{y}$, is expected to be normalized (i.e. `softmax` output). The modulating factor $\gamma$, controls the down-weighting strength. For $\gamma = 0$ this is equivalent to [`CrossEntropyLoss`](/api/Lux/utilities#Lux.CrossEntropyLoss). **Example** ```julia julia> y = [1 0 0 0 1 0 1 0 1 0 0 0 1 0 0] 3×5 Matrix{Int64}: 1 0 0 0 1 0 1 0 1 0 0 0 1 0 0 julia> ŷ = softmax(reshape(-7:7, 3, 5) .* 1f0) 3×5 Matrix{Float32}: 0.0900306 0.0900306 0.0900306 0.0900306 0.0900306 0.244728 0.244728 0.244728 0.244728 0.244728 0.665241 0.665241 0.665241 0.665241 0.665241 julia> FocalLoss()(ŷ, y) ≈ 1.1277556f0 true ``` source ```julia HingeLoss(; agg = mean) ``` Return the hinge loss loss given the prediction `ŷ` and true labels `y` (containing 1 or -1); calculated as: $$\text{agg}\left(\max(0, 1 - y \hat{y})\right)$$ Usually used with classifiers like Support Vector Machines. **Example** ```julia julia> loss = HingeLoss(); julia> y_true = [1, -1, 1, 1]; julia> y_pred = [0.1, 0.3, 1, 1.5]; julia> loss(y_pred, y_true) ≈ 0.55 true ``` source ```julia HuberLoss(; delta = 1, agg = mean) ``` Returns the Huber loss, calculated as: $$L = \begin{cases} 0.5 \* |y - \hat{y}|^2 & \text{if } |y - \hat{y}| \leq \delta \\ \delta \* (|y - \hat{y}| - 0.5 \* \delta) & \text{otherwise} \end{cases}$$ where $\delta$ is the `delta` parameter. **Example** ```julia julia> y_model = [1.1, 2.1, 3.1]; julia> HuberLoss()(y_model, 1:3) ≈ 0.005000000000000009 true julia> HuberLoss(delta=0.05)(y_model, 1:3) ≈ 0.003750000000000005 true ``` source ```julia KLDivergenceLoss(; dims = 1, agg = mean, epsilon = nothing, label_smoothing = nothing) ``` Return the Kullback-Leibler Divergence loss between the predicted distribution $\hat{y}$ and the true distribution $y$: The KL divergence is a measure of how much one probability distribution is different from the other. It is always non-negative, and zero only when both the distributions are equal. For `epsilon` and `label_smoothing`, see [`CrossEntropyLoss`](/api/Lux/utilities#Lux.CrossEntropyLoss). **Example** ```julia julia> p1 = [1 0; 0 1] 2×2 Matrix{Int64}: 1 0 0 1 julia> p2 = fill(0.5, 2, 2) 2×2 Matrix{Float64}: 0.5 0.5 0.5 0.5 julia> KLDivergenceLoss()(p2, p1) ≈ log(2) true julia> KLDivergenceLoss(; agg=sum)(p2, p1) ≈ 2 * log(2) true julia> KLDivergenceLoss(; epsilon=0)(p2, p2) 0.0 julia> KLDivergenceLoss(; epsilon=0)(p1, p2) Inf ``` source ```julia MAELoss(; agg = mean) ``` Returns the loss corresponding to mean absolute error: $$\text{agg}\left(\left| \hat{y} - y \right|\right)$$ **Example** ```julia julia> loss = MAELoss(); julia> y_model = [1.1, 1.9, 3.1]; julia> loss(y_model, 1:3) ≈ 0.1 true ``` source ```julia MSELoss(; agg = mean) ``` Returns the loss corresponding to mean squared error: $$\text{agg}\left(\left( \hat{y} - y \right)^2\right)$$ **Example** ```julia julia> loss = MSELoss(); julia> y_model = [1.1, 1.9, 3.1]; julia> loss(y_model, 1:3) ≈ 0.01 true ``` source ```julia MSLELoss(; agg = mean, epsilon = nothing) ``` Returns the loss corresponding to mean squared logarithmic error: $$\text{agg}\left(\left( \log\left( \hat{y} + \epsilon \right) - \log\left( y + \epsilon \right) \right)^2\right)$$ `epsilon` is added to both `y` and `ŷ` to prevent taking the logarithm of zero. If `epsilon` is `nothing`, then we set it to `eps()`. **Example** ```julia julia> loss = MSLELoss(); julia> loss(Float32[1.1, 2.2, 3.3], 1:3) ≈ 0.009084041f0 true julia> loss(Float32[0.9, 1.8, 2.7], 1:3) ≈ 0.011100831f0 true ``` source ```julia PoissonLoss(; agg = mean, epsilon = nothing) ``` Return how much the predicted distribution $\hat{y}$ diverges from the expected Poisson distribution $y$, calculated as: $$\text{agg}\left(\hat{y} - y \* \log(\hat{y})\right)$$ **Example** ```julia julia> y_model = [1, 3, 3]; # data should only take integral values julia> PoissonLoss()(y_model, 1:3) ≈ 0.502312852219817 true ``` source ```julia SiameseContrastiveLoss(; margin = true, agg = mean) ``` Return the contrastive loss \[[7](/references#hadsell2006dimensionality)] which can be useful for training Siamese Networks. It is given by: $$\text{agg}\left((1 - y) \hat{y}^2 + y \* \max(0, \text{margin} - \hat{y})^2\right)$$ Specify `margin` to set the baseline for distance at which pairs are dissimilar. **Example** ```julia julia> ŷ = [0.5, 1.5, 2.5]; julia> SiameseContrastiveLoss()(ŷ, 1:3) ≈ -4.833333333333333 true julia> SiameseContrastiveLoss(margin=2)(ŷ, 1:3) ≈ -4.0 true ``` source ```julia SquaredHingeLoss(; agg = mean) ``` Return the squared hinge loss loss given the prediction `ŷ` and true labels `y` (containing 1 or -1); calculated as: $$\text{agg}\left(\max(0, 1 - y \hat{y})^2\right)$$ Usually used with classifiers like Support Vector Machines. **Example** ```julia julia> loss = SquaredHingeLoss(); julia> y_true = [1, -1, 1, 1]; julia> y_pred = [0.1, 0.3, 1, 1.5]; julia> loss(y_pred, y_true) ≈ 0.625 true ``` source ## LuxOps Module {#LuxOps-Module} ```julia LuxOps ``` This module is a part of `Lux.jl`. It contains operations that are useful in DL context. Additionally certain operations here alias Base functions to behave more sensibly with GPUArrays. source ```julia eachslice(x, dims::Val) ``` Same as `Base.eachslice` but doesn't produce a `SubArray` for the slices if `x` is a GPUArray. Additional dispatches for RNN helpers are also provided for `TimeLastIndex` and `BatchLastIndex`. source ```julia foldl_init(op, x) foldl_init(op, x, init) ``` Exactly same as `foldl(op, x; init)` in the forward pass. But, gives gradients wrt `init` in the backward pass. source ```julia getproperty(x, ::Val{v}) getproperty(x, ::StaticSymbol{v}) ``` Similar to `Base.getproperty` but requires a `Val` (or `Static.StaticSymbol`). Additionally, if `v` is not present in `x`, then `nothing` is returned. source ```julia xlogx(x::Number) ``` Return `x * log(x)` for `x ≥ 0`, handling `x == 0` by taking the limit from above, to get zero. source ```julia xlogy(x::Number, y::Number) ``` Return `x * log(y)` for `y > 0`, and zero when `x == 0`. source ```julia istraining(::Val{training}) istraining(::StaticBool) istraining(::Bool) istraining(st::NamedTuple) ``` Returns `true` if `training` is `true` or if `st` contains a `training` field with value `true`. Else returns `false`. source ```julia multigate(x::AbstractArray, ::Val{N}) ``` Split up `x` into `N` equally sized chunks (along dimension `1`). source ```julia rsqrt(x) ``` Computes the reciprocal square root of `x`. For all backends except `Reactant.jl`, this falls back to `inv(sqrt(x))`. source ## Recursive Operations {#Recursive-Operations} ```julia recursive_map(f, x, args...) ``` Similar to `fmap(f, args...)` but with restricted support for the notion of "leaf" types. However, this allows for more efficient and type stable implementations of recursive operations. ::: warning Deprecation Warning Starting Lux v1.3.0, this function is deprecated in favor of `Functors.fmap`. Functors v0.5 made significant strides towards improving the performance of `fmap` and hence this function has been deprecated. Users are encouraged to use `Functors.fmap` instead. ::: **How this works?** For the following types it directly defines recursion rules: 1. `AbstractArray`: If eltype is `isbitstype`, then `f` is applied to the array, else we recurse on the array. 2. `Tuple/NamedTuple`: We recurse on the values. 3. `Number/Val/Nothing`: We directly apply `f`. 4. For all other types, we recurse on the fields using `Functors.fmap`. ::: tip Note In most cases, users should gravitate towards `Functors.fmap` if it is being used outside of hot loops. Even for other cases, it is always recommended to verify the correctness of this implementation for specific usecases. ::: source ```julia recursive_add!!(x, y) ``` Recursively add the leaves of two nested structures `x` and `y`. In Functor language, this is equivalent to doing `fmap(+, x, y)`, but this implementation uses type stable code for common cases. Any leaves of `x` that are arrays and allow in-place addition will be modified in place. ::: warning Deprecation Warning Starting Lux v1.3.0, this function is deprecated in favor of `Functors.fmap`. Functors v0.5 made significant strides towards improving the performance of `fmap` and hence this function has been deprecated. Users are encouraged to use `Functors.fmap` instead. ::: source ```julia recursive_copyto!(x, y) ``` Recursively copy the leaves of two nested structures `x` and `y`. In Functor language, this is equivalent to doing `fmap(copyto!, x, y)`, but this implementation uses type stable code for common cases. Note that any immutable leaf will lead to an error. ::: warning Deprecation Warning Starting Lux v1.3.0, this function is deprecated in favor of `Functors.fmap`. Functors v0.5 made significant strides towards improving the performance of `fmap` and hence this function has been deprecated. Users are encouraged to use `Functors.fmap` instead. ::: source ```julia recursive_eltype(x, unwrap_ad_types = Val(false)) ``` Recursively determine the element type of a nested structure `x`. This is equivalent to doing `fmap(Lux.Utils.eltype, x)`, but this implementation uses type stable code for common cases. For ambiguous inputs like `nothing` and `Val` types we return `Bool` as the eltype. If `unwrap_ad_types` is set to `Val(true)` then for tracing and operator overloading based ADs (ForwardDiff, ReverseDiff, Tracker), this function will return the eltype of the unwrapped value. source ```julia recursive_make_zero(x) ``` Recursively create a zero value for a nested structure `x`. This is equivalent to doing `fmap(zero, x)`, but this implementation uses type stable code for common cases. See also [`Lux.recursive_make_zero!!`](/api/Lux/utilities#Lux.recursive_make_zero!!). ::: warning Deprecation Warning Starting Lux v1.3.0, this function is deprecated in favor of `Functors.fmap`. Functors v0.5 made significant strides towards improving the performance of `fmap` and hence this function has been deprecated. Users are encouraged to use `Functors.fmap` instead. ::: source ```julia recursive_make_zero!!(x) ``` Recursively create a zero value for a nested structure `x`. Leaves that can be mutated with in-place zeroing will be modified in place. See also [`Lux.recursive_make_zero`](/api/Lux/utilities#Lux.recursive_make_zero) for fully out-of-place version. ::: warning Deprecation Warning Starting Lux v1.3.0, this function is deprecated in favor of `Functors.fmap`. Functors v0.5 made significant strides towards improving the performance of `fmap` and hence this function has been deprecated. Users are encouraged to use `Functors.fmap` instead. ::: source ## Updating Floating Point Precision {#Updating-Floating-Point-Precision} By default, Lux uses Float32 for all parameters and states. To update the precision simply pass the parameters / states / arrays into one of the following functions. ```julia f16(m) ``` Converts the `eltype` of `m` *floating point* values to `Float16`. To avoid recursion into structs mark them with `Functors.@leaf`. source ```julia f32(m) ``` Converts the `eltype` of `m` *floating point* values to `Float32`. To avoid recursion into structs mark them with `Functors.@leaf`. source ```julia f64(m) ``` Converts the `eltype` of `m` *floating point* values to `Float64`. To avoid recursion into structs mark them with `Functors.@leaf`. source ```julia bf16(m) ``` Converts the `eltype` of `m` *floating point* values to `BFloat16`. To avoid recursion into structs mark them with `Functors.@leaf`. ::: warning Warning `BFloat16s.jl` needs to be loaded before using this function. ::: ::: tip Support for `BFloat16` Most Lux operations aren't optimized for `BFloat16` yet. Instead this is meant to be used together with `Reactant.@compile`. ::: source ## Element Type Matching {#Element-Type-Matching} ```julia match_eltype(layer, ps, st, args...) ``` Helper function to "maybe" (see below) match the element type of `args...` with the element type of the layer's parameters and states. This is useful for debugging purposes, to track down accidental type-promotions inside Lux layers. **Extended Help** **Controlling the Behavior via Preferences** Behavior of this function is controlled via the `eltype_mismatch_handling` preference. The following options are supported: * `"none"`: This is the default behavior. In this case, this function is a no-op, i.e., it simply returns `args...`. * `"warn"`: This option will issue a warning if the element type of `args...` does not match the element type of the layer's parameters and states. The warning will contain information about the layer and the element type mismatch. * `"convert"`: This option is same as `"warn"`, but it will also convert the element type of `args...` to match the element type of the layer's parameters and states (for the cases listed below). * `"error"`: Same as `"warn"`, but instead of issuing a warning, it will throw an error. ::: warning Warning We print the warning for type-mismatch only once. ::: **Element Type Conversions** For `"convert"` only the following conversions are done: | Element Type of parameters/states | Element Type of `args...` | Converted to | |:--------------------------------- |:------------------------- |:------------ | | `Float64` | `Integer` | `Float64` | | `Float32` | `Float64` | `Float32` | | `Float32` | `Integer` | `Float32` | | `Float16` | `Float64` | `Float16` | | `Float16` | `Float32` | `Float16` | | `Float16` | `Integer` | `Float16` | source ## Stateful Layer {#Stateful-Layer} This layer have been moved to \[`LuxCore.jl`]. See the [documentation](/api/Building_Blocks/LuxCore#LuxCore.StatefulLuxLayerImpl.StatefulLuxLayer) for more details. ## Compact Layer {#Compact-Layer} ```julia @compact(kw...) do x ... @return y # optional (but recommended for best performance) end @compact(kw...) do x, p ... @return y # optional (but recommended for best performance) end @compact(forward::Function; name=nothing, dispatch=nothing, parameters...) ``` Creates a layer by specifying some `parameters`, in the form of keywords, and (usually as a `do` block) a function for the forward pass. You may think of `@compact` as a specialized `let` block creating local variables that are trainable in Lux. Declared variable names may be used within the body of the `forward` function. Note that unlike typical Lux models, the forward function doesn't need to explicitly manage states. Defining the version with `p` allows you to access the parameters in the forward pass. This is useful when using it with SciML tools which require passing in the parameters explicitly. **Reserved Kwargs:** 1. `name`: The name of the layer. 2. `dispatch`: The constructed layer has the type `Lux.CompactLuxLayer{dispatch}` which can be used for custom dispatches. ::: tip Tip Check the Lux tutorials for more examples of using `@compact`. ::: If you are passing in kwargs by splatting them, they will be passed as is to the function body. This means if your splatted kwargs contain a lux layer that won't be registered in the CompactLuxLayer. Additionally all of the device functions treat these kwargs as leaves. **Special Syntax** * `@return`: This macro doesn't really exist, but is used to return a value from the `@compact` block. Without the presence of this macro, we need to rely on closures which can lead to performance penalties in the reverse pass. * Having statements after the last `@return` macro might lead to incorrect code. * Don't do things like `@return return x`. This will generate non-sensical code like ` = return x`. Essentially, `@return ` supports any expression, that can be assigned to a variable. * Since this macro doesn't "exist", it cannot be imported as `using Lux: @return`. Simply use it in code, and `@compact` will understand it. * `@init_fn`: Provide a function that will be used to initialize the layer's parameters or state. See the docs of [`@init_fn`](/api/Lux/utilities#Lux.@init_fn) for more details. * `@non_trainable`: Mark a value as non-trainable. This bypasses the regular checks and places the value into the state of the layer. See the docs of [`@non_trainable`](/api/Lux/utilities#Lux.@non_trainable) for more details. **Extended Help** **Examples** Here is a linear model: ```julia julia> using Lux, Random julia> r = @compact(w=ones(Float32, 3)) do x @return w .* x end @compact( w = 3-element Vector{Float32}, ) do x return w .* x end # Total: 3 parameters, # plus 0 states. julia> ps, st = Lux.setup(Xoshiro(0), r); julia> r(Float32[1, 2, 3], ps, st) # x is set to [1, 1, 1]. (Float32[1.0, 2.0, 3.0], NamedTuple()) ``` Here is a linear model with bias and activation: ```julia julia> d_in = 5 5 julia> d_out = 3 3 julia> d = @compact(W=ones(Float32, d_out, d_in), b=zeros(Float32, d_out), act=relu) do x y = W * x @return act.(y .+ b) end @compact( W = 3×5 Matrix{Float32}, b = 3-element Vector{Float32}, act = relu, ) do x y = W * x return act.(y .+ b) end # Total: 18 parameters, # plus 1 states. julia> ps, st = Lux.setup(Xoshiro(0), d); julia> d(ones(Float32, 5, 2), ps, st)[1] # 3×2 Matrix as output. 3×2 Matrix{Float32}: 5.0 5.0 5.0 5.0 5.0 5.0 julia> ps_dense = (; weight=ps.W, bias=ps.b); julia> first(d(Float32[1, 2, 3, 4, 5], ps, st)) ≈ first(Dense(d_in => d_out, relu)(Float32[1, 2, 3, 4, 5], ps_dense, NamedTuple())) # Equivalent to a dense layer true ``` Finally, here is a simple MLP. We can train this model just like any Lux model: ```julia julia> n_in = 1; julia> n_out = 1; julia> nlayers = 3; julia> model = @compact(w1=Dense(n_in, 128), w2=[Dense(128, 128) for i in 1:nlayers], w3=Dense(128, n_out), act=relu) do x embed = act.(w1(x)) for w in w2 embed = act.(w(embed)) end out = w3(embed) @return out end @compact( w1 = Dense(1 => 128), # 256 parameters w2 = NamedTuple( (1-3) = Dense(128 => 128), # 49_536 (16_512 x 3) parameters ), w3 = Dense(128 => 1), # 129 parameters act = relu, ) do x embed = act.(w1(x)) for w = w2 embed = act.(w(embed)) end out = w3(embed) return out end # Total: 49_921 parameters, # plus 1 states. julia> ps, st = Lux.setup(Xoshiro(0), model); julia> size(first(model(randn(Float32, n_in, 32), ps, st))) # 1×32 Matrix as output. (1, 32) julia> using Optimisers, Zygote julia> x_data = collect(-2.0f0:0.1f0:2.0f0)'; julia> y_data = 2 .* x_data .- x_data .^ 3; julia> optim = Optimisers.setup(Adam(), ps); julia> loss_initial = sum(abs2, first(model(x_data, ps, st)) .- y_data); julia> for epoch in 1:1000 loss, gs = Zygote.withgradient( ps -> sum(abs2, first(model(x_data, ps, st)) .- y_data), ps) Optimisers.update!(optim, ps, gs[1]) end; julia> loss_final = sum(abs2, first(model(x_data, ps, st)) .- y_data); julia> loss_initial > loss_final true ``` You may also specify a `name` for the model, which will be used instead of the default printout, which gives a verbatim representation of the code used to construct the model: ```julia julia> model = @compact(w=rand(Float32, 3), name="Linear(3 => 1)") do x @return sum(w .* x) end Linear(3 => 1) # 3 parameters ``` This can be useful when using `@compact` to hierarchically construct complex models to be used inside a `Chain`. ::: tip Type Stability If your input function `f` is type-stable but the generated model is not type stable, it should be treated as a bug. We will appreciate issues if you find such cases. ::: ::: warning Parameter Count Array Parameter don't print the number of parameters on the side. However, they do account for the total number of parameters printed at the bottom. ::: source ```julia @init_fn(fn, kind::Symbol = :parameter) ``` Create an initializer function for a parameter or state to be used for in a Compact Lux Layer created using [`@compact`](/api/Lux/utilities#Lux.@compact). **Arguments** * `fn`: The function to be used for initializing the parameter or state. This only takes a single argument `rng`. * `kind`: If set to `:parameter`, the initializer function will be used to initialize the parameters of the layer. If set to `:state`, the initializer function will be used to initialize the states of the layer. **Examples** ```julia julia> using Lux, Random julia> r = @compact(w=@init_fn(rng->randn32(rng, 3, 2)), b=@init_fn(rng->randn32(rng, 3), :state)) do x @return w * x .+ b end; julia> ps, st = Lux.setup(Xoshiro(0), r); julia> size(ps.w) (3, 2) julia> size(st.b) (3,) julia> size(r([1, 2], ps, st)[1]) (3,) ``` source ```julia @non_trainable(x) ``` Mark a value as non-trainable. This bypasses the regular checks and places the value into the state of the layer. **Arguments** * `x`: The value to be marked as non-trainable. **Examples** ```julia julia> using Lux, Random julia> r = @compact(w=ones(3), w_fixed=@non_trainable(rand(3))) do x @return sum(w .* x .+ w_fixed) end; julia> ps, st = Lux.setup(Xoshiro(0), r); julia> size(ps.w) (3,) julia> size(st.w_fixed) (3,) julia> res, st_ = r([1, 2, 3], ps, st); julia> st_.w_fixed == st.w_fixed true julia> res isa Number true ``` source ## Miscellaneous {#Miscellaneous} ```julia set_dispatch_doctor_preferences!(mode::String) set_dispatch_doctor_preferences!(; luxcore::String="disable", luxlib::String="disable") ``` Set the dispatch doctor preference for `LuxCore` and `LuxLib` packages. `mode` can be `"disable"`, `"warn"`, or `"error"`. For details on the different modes, see the [DispatchDoctor.jl](https://astroautomata.com/DispatchDoctor.jl/dev/) documentation. If the preferences are already set, then no action is taken. Otherwise the preference is set. For changes to take effect, the Julia session must be restarted. source --- --- url: /dev/api/Lux/contrib.md --- # Experimental Features {#Experimental-Features} All features listed on this page are **experimental** which means: 1. No SemVer Guarantees. We use code here to iterate fast. That said, historically we have never broken any code in this module and have always provided a deprecation period. 2. Expect edge-cases and report them. It will help us move these features out of experimental sooner. 3. None of the features are exported. ## Parameter Freezing {#Parameter-Freezing} ```julia FrozenLayer(l::AbstractLuxLayer, which_params::Optional{Tuple}) ``` Freeze the parameters with name `which_params` of the layer `l`. ::: tip Use `Lux.Experimental.freeze` instead It is always recommended to use the [`Lux.Experimental.freeze`](/api/Lux/contrib#Lux.Experimental.freeze) function instead of directly using the `FrozenLayer` constructor. ::: ::: warning No checks for `which_params` There are no checks for `which_params`. For example, if the original layer has parameters named `(:weight, :bias)`, and `which_params` is set to `(:myweight,)` then none of the parameters are frozen and no error is thrown. ::: **Arguments** * `l`: Lux AbstractLuxLayer. * `which_params`: Parameter Names to be Frozen. Can be set to `nothing`, in which case all parameters are frozen. **Extended Help** **Parameters** * Parameters of the layer `l` excluding `which_params`. **States** * `frozen_params`: Parameters that are frozen, i.e., `which_params`. * `states`: The state of the inner layer `l`. **Note on Internal Layer Implementation** The inner layer should work with `NamedTuple` parameters. In order to support custom parameter types, users need to implement `Lux.Utils.merge(::CustomParamType, ::NamedTuple)` or extend `Lux.Utils.named_tuple(::CustomParamType)` to return a `NamedTuple`. **Example** ```julia julia> Lux.Experimental.FrozenLayer(Dense(2 => 2), (:weight,)) FrozenLayer(Dense(2 => 2), (:weight,)) # 2 parameters, plus 4 non-trainable ``` See also [`Lux.Experimental.freeze`](/api/Lux/contrib#Lux.Experimental.freeze), [`Lux.Experimental.unfreeze`](/api/Lux/contrib#Lux.Experimental.unfreeze). source ```julia freeze(l::AbstractLuxLayer, which_params::Optional{Tuple} = nothing) ``` Constructs a version of `l` with `which_params` frozen. If `which_params` is nothing, then all parameters are frozen. source ```julia freeze(l::AbstractLuxLayer, ps, st::NamedTuple, which_params::Optional{Tuple} = nothing) ``` Construct a [`Lux.Experimental.FrozenLayer`](/api/Lux/contrib#Lux.Experimental.FrozenLayer) for `l` with the current parameters and states. If `which_params` is nothing, then all parameters are frozen. source ```julia unfreeze(l::FrozenLayer) ``` Unfreezes the layer `l`. source ```julia unfreeze(l::FrozenLayer, ps, st::NamedTuple) ``` Unwraps a [`Lux.Experimental.FrozenLayer`](/api/Lux/contrib#Lux.Experimental.FrozenLayer) `l` with the current parameters and states. source For detailed usage example look at the [manual page](/manual/freezing_model_parameters#freezing-model-parameters). ## Map over Layer {#Map-over-Layer} ```julia layer_map(f, l::AbstractLuxLayer, ps, st::NamedTuple) ``` Map the function `f` over the model `l`, with the parameters `ps` and states `st`. This is different from `Functors.fmap` since it zips the layers, parameters, and states and invokes the function on all of them together. ::: tip KeyPath provided to the function The `KeyPath` depths on the structure of the parameters and states. This is of consequence exclusively for [`AbstractLuxWrapperLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxWrapperLayer) where the structure of the layer doesn't match the structure of the parameters and states. In the example, provided below, the `KeyPath` is `(:chain, :dense_1)` for the first layer (following the structure in `ps`) while accessing the same layer in the chain is done with `( :chain, :layers, :dense_1)`. ::: **Call Signature for `f`** * Must take 4 inputs – `AbstractLuxLayer`, Corresponding Parameters, Corresponding States, and the `Functors.KeyPath` to the layer. * Must return a tuple of 3 elements – `AbstractLuxLayer`, new parameters and the new states. **Extended Help** **Example** ```julia julia> using Lux, Random julia> c = Parallel( +; chain=Chain(; dense_1=Dense(2 => 3), bn=BatchNorm(3), dense_2=Dense(3 => 5)), dense_3=Dense(5 => 1)); julia> rng = Random.default_rng(); julia> ps, st = Lux.setup(rng, c); julia> # Makes parameters of Dense Layers inside Chain zero function zero_dense_params(l, ps, st, name) if l isa Dense println("zeroing params of $name") ps = merge(ps, (; weight=zero.(ps.weight), bias=zero.(ps.bias))) end return l, ps, st end; julia> _, ps_new, _ = Lux.Experimental.layer_map(zero_dense_params, c, ps, st); zeroing params of KeyPath(:chain, :dense_1) zeroing params of KeyPath(:chain, :dense_2) zeroing params of KeyPath(:dense_3,) julia> all(iszero, (ps_new.chain.dense_1.weight, ps_new.chain.dense_1.bias, ps_new.chain.dense_2.weight, ps_new.chain.dense_2.bias, ps_new.dense_3.weight, ps_new.dense_3.bias)) true ``` source ## Debugging Functionality {#Debugging-Functionality} Model not working properly! Here are some functionalities to help you debug you Lux model. ```julia @debug_mode layer kwargs... ``` Recurses into the `layer` and replaces the inner most non Container Layers with a [`Lux.Experimental.DebugLayer`](/api/Lux/contrib#Lux.Experimental.DebugLayer). See [`Lux.Experimental.DebugLayer`](/api/Lux/contrib#Lux.Experimental.DebugLayer) for details about the Keyword Arguments. source ```julia DebugLayer(layer::AbstractLuxLayer; nan_check::Union{Symbol, StaticSymbol, Val}=static(:both), error_check::Union{StaticBool, Bool, Val{true}, Val{false}}=True(), location::KeyPath=KeyPath()) ``` A wrapper over Lux layers that adds checks for NaNs and errors. This is useful for debugging. **Arguments** * `layer`: The layer to be wrapped. **Extended Help** **Keyword Arguments** * `nan_check`: Whether to check for NaNs in the input, parameters, and states. Can be `:both`, `:forward`, `:backward`, or `:none`. * `error_check`: Whether to check for errors in the layer. If `true`, will throw an error if the layer fails. * `location`: The location of the layer. Use [`Lux.Experimental.@debug_mode`](/api/Lux/contrib#Lux.Experimental.@debug_mode) to construct this layer to populate this value correctly. **Input / Output** Inputs and outputs are the same as the `layer` unless one of the `nan_check` or `error_check` criteria is met. If `nan_check` is enabled and NaNs are detected then a `DomainError` is thrown. If `error_check` is enabled, then any errors in the layer are thrown with useful information to track where the error originates. ::: warning ChainRules Compatible Reverse Mode AD Tools `nan_check` for the backward mode only works with ChainRules Compatible Reverse Mode AD Tools currently. ::: ::: danger Disable After Debugging This layer is only meant to be used for debugging. If used for actual training or inference, will lead to extremely bad performance. ::: See [`Lux.Experimental.@debug_mode`](/api/Lux/contrib#Lux.Experimental.@debug_mode) to construct this layer. source ## Tied Parameters {#Tied-Parameters} ```julia share_parameters(ps, sharing) share_parameters(ps, sharing, new_parameters) ``` Updates the parameters in `ps` with a common set of parameters `new_parameters` that are shared between each list in the nested list `sharing`. (That was kind of a mouthful, the example should make it clear). **Arguments** * `ps`: Original parameters. * `sharing`: A nested list of lists of accessors of `ps` which need to shate the parameters (See the example for details). (Each list in the list must be disjoint) * `new_parameters`: If passed the length of `new_parameters` must be equal to the length of `sharing`. For each vector in `sharing` the corresponding parameter in `new_parameters` will be used. (If not passed, the parameters corresponding to the first element of each vector in `sharing` will be used). **Returns** Updated Parameters having the same structure as `ps`. **Example** ```julia julia> model = Chain(; d1=Dense(2 => 4, tanh), d3=Chain(; l1=Dense(4 => 2), l2=Dense(2 => 4)), d2=Dense(4 => 2)) Chain( d1 = Dense(2 => 4, tanh), # 12 parameters d3 = Chain( l1 = Dense(4 => 2), # 10 parameters l2 = Dense(2 => 4), # 12 parameters ), d2 = Dense(4 => 2), # 10 parameters ) # Total: 44 parameters, # plus 0 states. julia> ps, st = Lux.setup(Xoshiro(0), model); julia> # share parameters of (d1 and d3.l1) and (d3.l2 and d2) ps = Lux.Experimental.share_parameters(ps, (("d3.l2", "d1"), ("d2", "d3.l1"))); julia> ps.d3.l2.weight === ps.d1.weight && ps.d3.l2.bias === ps.d1.bias && ps.d2.weight === ps.d3.l1.weight && ps.d2.bias === ps.d3.l1.bias true ``` ::: danger ComponentArrays ComponentArrays doesn't allow sharing parameters. Converting the returned parameters to a ComponentArray will silently cause the parameter sharing to be undone. ::: source --- --- url: /dev/api/Lux/interop.md --- # Interoperability between Lux and other packages {#Interoperability-between-Lux-and-other-packages} ## Switching from older frameworks {#Switching-from-older-frameworks} ### Flux Models to Lux Models {#flux-to-lux-migrate-api} `Flux.jl` has been around in the Julia ecosystem for a long time and has a large userbase, hence we provide a way to convert `Flux` models to `Lux` models. ::: tip Tip Accessing these functions require manually loading `Flux`, i.e., `using Flux` must be present somewhere in the code for these to be used. ::: ```julia Adapt.adapt(from::FromFluxAdaptor, L) ``` Adapt a Flux model `L` to Lux model. See [`FromFluxAdaptor`](/api/Lux/interop#Lux.FromFluxAdaptor) for more details. source ```julia FromFluxAdaptor(preserve_ps_st::Bool=false, force_preserve::Bool=false) ``` Convert a Flux Model to Lux Model. ::: warning `active` field This always ignores the `active` field of some of the Flux layers. This is almost never going to be supported. ::: **Keyword Arguments** * `preserve_ps_st`: Set to `true` to preserve the states and parameters of the layer. This attempts the best possible way to preserve the original model. But it might fail. If you need to override possible failures, set `force_preserve` to `true`. * `force_preserve`: Some of the transformations with state and parameters preservation haven't been implemented yet, in these cases, if `force_transform` is `false` a warning will be printed and a core Lux layer will be returned. Else, it will create a [`FluxLayer`](/api/Lux/interop#Lux.FluxLayer). **Example** ```julia julia> import Flux julia> using Adapt, Lux, Random julia> m = Flux.Chain(Flux.Dense(2 => 3, relu), Flux.Dense(3 => 2)); julia> m2 = adapt(FromFluxAdaptor(), m); # or FromFluxAdaptor()(m.layers) julia> x = randn(Float32, 2, 32); julia> ps, st = Lux.setup(Random.default_rng(), m2); julia> size(first(m2(x, ps, st))) (2, 32) ``` source ```julia FluxLayer(layer) ``` Serves as a compatibility layer between Flux and Lux. This uses `Optimisers.destructure` API internally. ::: warning Warning Lux was written to overcome the limitations of `destructure` + `Flux`. It is recommended to rewrite your layer in Lux instead of using this layer. ::: ::: warning Warning Introducing this Layer in your model will lead to type instabilities, given the way `Optimisers.destructure` works. ::: **Arguments** * `layer`: Flux layer **Parameters** * `p`: Flattened parameters of the `layer` source ## Using a different backend for Lux {#Using-a-different-backend-for-Lux} ### Lux Models to Simple Chains {#Lux-Models-to-Simple-Chains} `SimpleChains.jl` provides a way to train Small Neural Networks really fast on CPUs. See [this blog post](https://julialang.org/blog/2022/04/simple-chains/) for more details. This section describes how to convert `Lux` models to `SimpleChains` models while preserving the [layer interface](/manual/interface#lux-interface). ::: tip Tip Accessing these functions require manually loading `SimpleChains`, i.e., `using SimpleChains` must be present somewhere in the code for these to be used. ::: ```julia Adapt.adapt(from::ToSimpleChainsAdaptor, L::AbstractLuxLayer) ``` Adapt a Simple Chains model to Lux model. See [`ToSimpleChainsAdaptor`](/api/Lux/interop#Lux.ToSimpleChainsAdaptor) for more details. source ```julia ToSimpleChainsAdaptor(input_dims, convert_to_array::Bool=false) ``` Adaptor for converting a Lux Model to SimpleChains. The returned model is still a Lux model, and satisfies the `AbstractLuxLayer` interfacem but all internal calculations are performed using SimpleChains. ::: warning Warning There is no way to preserve trained parameters and states when converting to `SimpleChains.jl`. ::: ::: warning Warning Any kind of initialization function is not preserved when converting to `SimpleChains.jl`. ::: **Arguments** * `input_dims`: Tuple of input dimensions excluding the batch dimension. These must be of `static` type as `SimpleChains` expects. * `convert_to_array`: SimpleChains.jl by default outputs `StrideArraysCore.StrideArray`, but this might not compose well with other packages. If `convert_to_array` is set to `true`, the output will be converted to a regular `Array`. **Example** ```julia julia> import SimpleChains julia> using Adapt, Lux, Random julia> lux_model = Chain(Conv((5, 5), 1 => 6, relu), MaxPool((2, 2)), Conv((5, 5), 6 => 16, relu), MaxPool((2, 2)), FlattenLayer(3), Chain(Dense(256 => 128, relu), Dense(128 => 84, relu), Dense(84 => 10))); julia> adaptor = ToSimpleChainsAdaptor((28, 28, 1)); julia> simple_chains_model = adapt(adaptor, lux_model); # or adaptor(lux_model) julia> ps, st = Lux.setup(Random.default_rng(), simple_chains_model); julia> x = randn(Float32, 28, 28, 1, 1); julia> size(first(simple_chains_model(x, ps, st))) (10, 1) ``` source ```julia SimpleChainsLayer(layer, to_array::Union{Bool, Val}=Val(false)) SimpleChainsLayer(layer, lux_layer, to_array) ``` Wraps a `SimpleChains` layer into a `Lux` layer. All operations are performed using `SimpleChains` but the layer satisfies the `AbstractLuxLayer` interface. `ToArray` is a boolean flag that determines whether the output should be converted to a regular `Array` or not. Default is `false`. **Arguments** * `layer`: SimpleChains layer * `lux_layer`: Potentially equivalent Lux layer that is used for printing source --- --- url: /dev/api/Lux/distributed_utils.md --- # Distributed Utils {#Distributed-Utils} ::: tip Note These functionalities are available via the `Lux.DistributedUtils` module. ::: ## Backends {#communication-backends} ```julia MPIBackend(comm = nothing) ``` Create an MPI backend for distributed training. Users should not use this function directly. Instead use [`DistributedUtils.get_distributed_backend(MPIBackend)`](/api/Lux/distributed_utils#Lux.DistributedUtils.get_distributed_backend). source ```julia NCCLBackend(comm = nothing, mpi_backend = nothing) ``` Create an NCCL backend for distributed training. Users should not use this function directly. Instead use [`DistributedUtils.get_distributed_backend(NCCLBackend)`](/api/Lux/distributed_utils#Lux.DistributedUtils.get_distributed_backend). source ## Initialization {#Initialization} ```julia initialize(backend::Type{<:AbstractLuxDistributedBackend}; kwargs...) ``` Initialize the given backend. Users can supply `cuda_devices` and `amdgpu_devices` to initialize the backend with the given devices. These can be set to `missing` to prevent initialization of the given device type. If set to `nothing`, and the backend is functional we assign GPUs in a round-robin fashion. Finally, a list of integers can be supplied to initialize the backend with the given devices. Possible values for `backend` are: * `MPIBackend`: MPI backend for distributed training. Requires `MPI.jl` to be installed. * `NCCLBackend`: NCCL backend for CUDA distributed training. Requires `CUDA.jl`, `MPI.jl`, and `NCCL.jl` to be installed. This also wraps `MPI` backend for non-CUDA communications. source ```julia initialized(backend::Type{<:AbstractLuxDistributedBackend}) ``` Check if the given backend is initialized. source ```julia get_distributed_backend(backend::Type{<:AbstractLuxDistributedBackend}) ``` Get the distributed backend for the given backend type. Possible values are: * `MPIBackend`: MPI backend for distributed training. Requires `MPI.jl` to be installed. * `NCCLBackend`: NCCL backend for CUDA distributed training. Requires `CUDA.jl`, `MPI.jl`, and `NCCL.jl` to be installed. This also wraps `MPI` backend for non-CUDA communications. ::: danger Danger `initialize(backend; kwargs...)` must be called before calling this function. ::: source ## Helper Functions {#Helper-Functions} ```julia local_rank(backend::AbstractLuxDistributedBackend) ``` Get the local rank for the given backend. source ```julia total_workers(backend::AbstractLuxDistributedBackend) ``` Get the total number of workers for the given backend. source ## Communication Primitives {#Communication-Primitives} ```julia allreduce!(backend::AbstractLuxDistributedBackend, sendrecvbuf, op) allreduce!(backend::AbstractLuxDistributedBackend, sendbuf, recvbuf, op) ``` Backend Agnostic API to perform an allreduce operation on the given buffer `sendrecvbuf` or `sendbuf` and store the result in `recvbuf`. `op` allows a special `DistributedUtils.avg` operation that averages the result across all workers. source ```julia bcast!(backend::AbstractLuxDistributedBackend, sendrecvbuf; root::Int=0) bcast!(backend::AbstractLuxDistributedBackend, sendbuf, recvbuf; root::Int=0) ``` Backend Agnostic API to broadcast the given buffer `sendrecvbuf` or `sendbuf` to all workers into `recvbuf`. The value at `root` will be broadcasted to all other workers. source ```julia reduce!(backend::AbstractLuxDistributedBackend, sendrecvbuf, op; root::Int=0) reduce!(backend::AbstractLuxDistributedBackend, sendbuf, recvbuf, op; root::Int=0) ``` Backend Agnostic API to perform a reduce operation on the given buffer `sendrecvbuf` or `sendbuf` and store the result in `recvbuf`. `op` allows a special `DistributedUtils.avg` operation that averages the result across all workers. source ```julia synchronize!!(backend::AbstractLuxDistributedBackend, ps; root::Int=0) ``` Synchronize the given structure `ps` using the given backend. The value at `root` will be broadcasted to all other workers. source ## Optimizers.jl Integration {#Optimizers.jl-Integration} ```julia DistributedOptimizer(backend::AbstractLuxDistributedBacked, optimizer) ``` Wrap the `optimizer` in a `DistributedOptimizer`. Before updating the parameters, this averages the gradients across the processes using Allreduce. **Arguments** * `optimizer`: An Optimizer compatible with the Optimisers.jl package source ## MLUtils.jl Integration {#MLUtils.jl-Integration} ```julia DistributedDataContainer(backend::AbstractLuxDistributedBackend, data) ``` `data` must be compatible with `MLUtils` interface. The returned container is compatible with `MLUtils` interface and is used to partition the dataset across the available processes. ::: danger Load `MLUtils.jl` `MLUtils.jl` must be installed and loaded before using this. ::: source --- --- url: /dev/api/Lux/serialization.md --- # Serialization {#Serialization} ## TensorFlow SavedModel {#TensorFlow-SavedModel} ```julia export_as_tf_saved_model( model_dir::String, model::AbstractLuxLayer, x, ps, st; mode=:inference, force::Bool=false, ) ``` Serializes a Lux model to a TensorFlow SavedModel format. A SavedModel contains a complete TensorFlow program, including trained parameters (i.e, tf.Variables) and computation. It does not require the original model building code to run, which makes it useful for sharing or deploying with [TFLite](https://tensorflow.org/lite), [TensorFlow.js](https://js.tensorflow.org/), [TensorFlow Serving](https://www.tensorflow.org/tfx/serving/tutorials/Serving_REST_simple), or [TensorFlow Hub](https://tensorflow.org/hub). Refer to the [official documentation](https://www.tensorflow.org/guide/saved_model) for more details. ::: warning Load `Reactant.jl` and `PythonCall.jl` before using this function This function requires the `Reactant` and `PythonCall` extensions to be loaded. If you haven't done so, please load them before calling this function. ::: ::: tip All inputs must be on `reactant_device()` The inputs `x`, `ps`, and `st` must be on the device returned by `reactant_device()`. If you are using a GPU, ensure that the inputs are on the GPU device. ::: ::: danger Running the saved model Currently we don't support saving a dynamically shaped tensor. Hence, for inference the input must be the same shape as the one used during export. ::: ::: warning Transposed Inputs When providing inputs to the loaded model, ensure that the input tensors are transposed, i.e. if the inputs was `[S₁, S₂, ..., Sₙ]` during export, then the input to the loaded model should be `[Sₙ, ..., S₂, S₁]`. ::: **Arguments** * `model_dir`: The directory where the model will be saved. * `model`: The model to be saved. * `x`: The input to the model. * `ps`: The parameters of the model. * `st`: The states of the model. **Keyword Arguments** * `mode`: The mode of the model. Can be either `:inference` or `:training`. Defaults to `:inference`. If set to `:training`, we will call [`LuxCore.trainmode`](/api/Building_Blocks/LuxCore#LuxCore.trainmode) on the model state, else we will call [`LuxCore.testmode`](/api/Building_Blocks/LuxCore#LuxCore.testmode). * `force`: If `true`, the function will overwrite existing files in the specified directory. Defaults to `false`. If the directory is not empty and `force` is `false`, the function will throw an error. **Example** Export the model to a TensorFlow SavedModel format. ```julia using Lux, Reactant, PythonCall, Random dev = reactant_device() model = Chain( Conv((5, 5), 1 => 6, relu), BatchNorm(6), MaxPool((2, 2)), Conv((5, 5), 6 => 16, relu), BatchNorm(16), MaxPool((2, 2)), FlattenLayer(3), Chain(Dense(256 => 128, relu), Dense(128 => 84, relu), Dense(84 => 10)), ) rng = Random.default_rng() ps, st = Lux.setup(rng, model) |> dev; x = rand(Float32, 28, 28, 1, 4) |> dev; Lux.Serialization.export_as_tf_saved_model("/tmp/testing_tf_saved_model", model, x, ps, st) ``` Load the model and run inference on a random input tensor. ```python import tensorflow as tf import numpy as np x_tf = tf.constant(np.random.rand(4, 1, 28, 28), dtype=tf.float32) restored_model = tf.saved_model.load("/tmp/testing_tf_saved_model") restored_model.f(x_tf)[0] ``` source --- --- url: /dev/api/Accelerator_Support/MLDataDevices.md --- # MLDataDevices {#MLDataDevices-API} `MLDataDevices.jl` is a lightweight package defining rules for transferring data across devices. ## Preferences {#Preferences} ```julia gpu_backend!() = gpu_backend!("") gpu_backend!(backend) = gpu_backend!(string(backend)) gpu_backend!(backend::AbstractGPUDevice) gpu_backend!(backend::String) ``` Creates a `LocalPreferences.toml` file with the desired GPU backend. If `backend == ""`, then the `gpu_backend` preference is deleted. Otherwise, `backend` is validated to be one of the possible backends and the preference is set to `backend`. If a new backend is successfully set, then the Julia session must be restarted for the change to take effect. source ## Data Transfer {#Data-Transfer} ```julia cpu_device(eltype=missing) -> CPUDevice ``` Return a `CPUDevice` object which can be used to transfer data to CPU. The `eltype` parameter controls element type conversion: * `missing/nothing` (default): Preserves the original element type * `Type{<:AbstractFloat}`: Converts floating-point arrays to the specified type source ```julia gpu_device( eltype::Union{Missing, Nothing, Type{<:AbstractFloat}}=missing; kwargs... ) -> AbstractDevice gpu_device( device_id::Union{Nothing, Integer}=nothing eltype::Union{Missing, Nothing, Type{<:AbstractFloat}}=missing; force::Bool=false ) -> AbstractDevice ``` Selects GPU device based on the following criteria: 1. If `gpu_backend` preference is set and the backend is functional on the system, then that device is selected. 2. Otherwise, an automatic selection algorithm is used. We go over possible device backends in the order specified by `supported_gpu_backends()` and select the first functional backend. 3. If no GPU device is functional and `force` is `false`, then `cpu_device()` is invoked. 4. If nothing works, an error is thrown. **Arguments** * `device_id::Union{Nothing, Integer}`: The device id to select. If `nothing`, then we return the last selected device or if none was selected then we run the autoselection and choose the current device using `CUDA.device()` or `AMDGPU.device()` or similar. If `Integer`, then we select the device with the given id. Note that this is `1`-indexed, in contrast to the `0`-indexed `CUDA.jl`. For example, `id = 4` corresponds to `CUDA.device!(3)`. * `eltype::Union{Missing, Nothing, Type{<:AbstractFloat}}`: The element type to use for the device. * `missing` (default): Device specific. For `CUDADevice` this calls `CUDA.cu(x)`, for `AMDGPUDevice` this calls `AMDGPU.roc(x)`, for `MetalDevice` this calls `Metal.mtl(x)`, for `oneAPIDevice` this calls `oneArray(x)`, for `OpenCLDevice` this calls `CLArray(x)`. * `nothing`: Preserves the original element type. * `Type{<:AbstractFloat}`: Converts floating-point arrays to the specified type. ::: warning Warning `device_id` is only applicable for `CUDA` and `AMDGPU` backends. For `Metal`, `oneAPI`, `OpenCL` and `CPU` backends, `device_id` is ignored and a warning is printed. ::: ::: warning Warning `gpu_device` won't select a CUDA device unless both CUDA.jl and cuDNN.jl are loaded. This is to ensure that deep learning operations work correctly. Nonetheless, if cuDNN is not loaded you can still manually create a `CUDADevice` object and use it (e.g. `dev = CUDADevice()`). ::: **Keyword Arguments** * `force::Bool`: If `true`, then an error is thrown if no functional GPU device is found. source ```julia reactant_device(; force::Bool=false, client=missing, device=missing, sharding=missing, eltype=missing, track_numbers::Type{TN}=Union{} ) -> Union{ReactantDevice, CPUDevice} ``` Return a `ReactantDevice` object if functional. Otherwise, throw an error if `force` is `true`. Falls back to `CPUDevice` if `force` is `false`. `client` and `device` are used to specify the client and particular device to use. If not specified, then the default client and index are used. `sharding` is used to specify the sharding strategy. If a `Reactant.Sharding.AbstractSharding` is specified, then we use it to shard all abstract arrays. Alternatively, pass in a `IdDict` to specify the sharding for specific leaves. `track_numbers` can be specified to convert numbers of specified subtypes to be traced. The `eltype` parameter controls element type conversion: * `missing/nothing` (default): Preserves the original element type * `Type{<:AbstractFloat}`: Converts floating-point arrays to the specified type source ## Miscellaneous {#Miscellaneous} ```julia reset_gpu_device!() ``` Resets the selected GPU device. This is useful when automatic GPU selection needs to be run again. source ```julia supported_gpu_backends() -> Tuple{String, ...} ``` Return a tuple of supported GPU backends. ::: warning Warning This is not the list of functional backends on the system, but rather backends which `MLDataDevices.jl` supports. ::: source ```julia default_device_rng(::AbstractDevice) ``` Returns the default RNG for the device. This can be used to directly generate parameters and states on the device using [WeightInitializers.jl](https://github.com/LuxDL/WeightInitializers.jl). source ```julia get_device(x) -> dev::AbstractDevice | Exception | Nothing ``` If all arrays (on the leaves of the structure) are on the same device, we return that device. Otherwise, we throw an error. If the object is device agnostic, we return `nothing`. ::: tip Note Trigger Packages must be loaded for this to return the correct device. ::: **Special Retuened Values** * `nothing` – denotes that the object is device agnostic. For example, scalar, abstract range, etc. * `UnknownDevice()` – denotes that the device type is unknown. See also [`get_device_type`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.get_device_type) for a faster alternative that can be used for dispatch based on device type. source ```julia get_device_type(x) -> Type{<:AbstractDevice} | Exception | Type{Nothing} ``` Similar to [`get_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.get_device) but returns the type of the device instead of the device itself. This value is often a compile time constant and is recommended to be used instead of [`get_device`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.get_device) where ever defining dispatches based on the device type. ::: tip Note Trigger Packages must be loaded for this to return the correct device. ::: **Special Retuened Values** * `Nothing` – denotes that the object is device agnostic. For example, scalar, abstract range, etc. * `UnknownDevice` – denotes that the device type is unknown. source ```julia loaded(x::AbstractDevice) -> Bool loaded(::Type{<:AbstractDevice}) -> Bool ``` Checks if the trigger package for the device is loaded. Trigger packages are as follows: * `CUDA.jl` and `cuDNN.jl` (or just `LuxCUDA.jl`) for NVIDIA CUDA Support. * `AMDGPU.jl` for AMD GPU ROCM Support. * `Metal.jl` for Apple Metal GPU Support. * `oneAPI.jl` for Intel oneAPI GPU Support. * `OpenCL.jl` for OpenCL support. source ```julia functional(x::AbstractDevice) -> Bool functional(::Type{<:AbstractDevice}) -> Bool ``` Checks if the device is functional. This is used to determine if the device can be used for computation. Note that even if the backend is loaded (as checked via [`MLDataDevices.loaded`](/api/Accelerator_Support/MLDataDevices#MLDataDevices.loaded)), the device may not be functional. Note that while this function is not exported, it is considered part of the public API. source ```julia isleaf(x) -> Bool ``` Returns `true` if `x` is a leaf node in the data structure. Defining `MLDataDevices.isleaf(x::T) = true` for custom types can be used to customize the behavior the data movement behavior when an object with nested structure containing the type is transferred to a device. `Adapt.adapt_structure(::AbstractDevice, x::T)` or `Adapt.adapt_structure(::AbstractDevice, x::T)` will be called during data movement if `isleaf(x::T)`. If `MLDataDevices.isleaf(x::T)` is not defined, then it will fall back to `Functors.isleaf(x)`. source ## Multi-GPU Support {#Multi-GPU-Support} ```julia set_device!(T::Type{<:AbstractDevice}, dev_or_id) ``` Set the device for the given type. This is a no-op for `CPUDevice`. For `CUDADevice` and `AMDGPUDevice`, it prints a warning if the corresponding trigger package is not loaded. Currently, `MetalDevice` and `oneAPIDevice` don't support setting the device. **Arguments** * `T::Type{<:AbstractDevice}`: The device type to set. * `dev_or_id`: Can be the device from the corresponding package. For example for CUDA it can be a `CuDevice`. If it is an integer, it is the device id to set. This is `1`-indexed. ::: danger Danger This specific function should be considered experimental at this point and is currently provided to support distributed training in Lux. As such please use `Lux.DistributedUtils` instead of using this function. ::: source ```julia set_device!(T::Type{<:AbstractDevice}, ::Nothing, rank::Integer) ``` Set the device for the given type. This is a no-op for `CPUDevice`. For `CUDADevice` and `AMDGPUDevice`, it prints a warning if the corresponding trigger package is not loaded. Currently, `MetalDevice` and `oneAPIDevice` don't support setting the device. **Arguments** * `T::Type{<:AbstractDevice}`: The device type to set. * `rank::Integer`: Local Rank of the process. This is applicable for distributed training and must be `0`-indexed. ::: danger Danger This specific function should be considered experimental at this point and is currently provided to support distributed training in Lux. As such please use `Lux.DistributedUtils` instead of using this function. ::: source ## Iteration {#Iteration} ```julia DeviceIterator(dev::AbstractDevice, iterator) ``` Create a `DeviceIterator` that iterates through the provided `iterator` via `iterate`. Upon each iteration, the current batch is copied to the device `dev`, and the previous iteration is marked as freeable from GPU memory (via `unsafe_free!`) (no-op for a CPU device). The conversion follows the same semantics as `dev()`. ::: tip Similarity to `CUDA.CuIterator` The design inspiration was taken from `CUDA.CuIterator` and was generalized to work with other backends and more complex iterators (using `Functors`). ::: ::: tip `MLUtils.DataLoader` Calling `dev(::MLUtils.DataLoader)` will automatically convert the dataloader to use the same semantics as `DeviceIterator`. This is generally preferred over looping over the dataloader directly and transferring the data to the device. ::: **Examples** The following was run on a computer with an NVIDIA GPU. ```julia julia> using MLDataDevices, MLUtils julia> X = rand(Float64, 3, 33); julia> dataloader = DataLoader(X; batchsize=13, shuffle=false); julia> for (i, x) in enumerate(dataloader) @show i, summary(x) end (i, summary(x)) = (1, "3×13 Matrix{Float64}") (i, summary(x)) = (2, "3×13 Matrix{Float64}") (i, summary(x)) = (3, "3×7 Matrix{Float64}") julia> for (i, x) in enumerate(CUDADevice()(dataloader)) @show i, summary(x) end (i, summary(x)) = (1, "3×13 CuArray{Float32, 2, CUDA.DeviceMemory}") (i, summary(x)) = (2, "3×13 CuArray{Float32, 2, CUDA.DeviceMemory}") (i, summary(x)) = (3, "3×7 CuArray{Float32, 2, CUDA.DeviceMemory}") ``` source --- --- url: /dev/api/NN_Primitives/LuxLib.md --- # LuxLib {#LuxLib-API} Backend for Lux.jl ## Apply Activation {#Apply-Activation} ```julia fast_activation(σ::F, x::AbstractArray) where {F} ``` Compute `σ.(x)` with the best possible implementation available. On CPUs we unroll the loop and use LoopVectorization.jl to vectorize the computation. On GPUs we use simply use broadcasting. ::: tip Note This function doesn't replace `σ` with `NNlib.fast_act(σ, ...)`, that needs to be done by the user if needed. ::: **Arguments** * `σ`: Activation function * `x`: Input array **Returns** * Output Array with the same size as `x` source ```julia fast_activation!!(σ::F, x::AbstractArray) where {F} ``` Compute `σ.(x)` with the best possible implementation available. If it is possible to rewrite `x` in-place, it does so. If `x` is an immutable array, it falls back to the generic implementation. ::: tip Note This function doesn't replace `σ` with `NNlib.fast_act(σ, ...)`, that needs to be done by the user if needed. ::: ::: tip Load `SLEEFPirates.jl` to get faster activations Certain activation functions are replaced with specialized implementations from [SLEEFPirates.jl](https://github.com/JuliaSIMD/SLEEFPirates.jl) for FP32. This might lead to faster performance but can cause slight decrease in accuracy (in the floating point limit). ::: **Arguments** * `σ`: Activation function * `x`: Input array **Returns** * Output Array with the same size as `x` source ## Attention {#Attention} ```julia scaled_dot_product_attention( query, key, value; head_dim=1, token_dim=3, scale=nothing, mask=nothing, is_causal=nothing, bias=nothing, fdrop=identity ) ``` Compute scaled dot-product attention for general tensors. This function is a thin API wrapper that forwards to the implementation in `impl/attention.jl` which supports flexible placement of the `head` and `token` dimensions (via `head_dim` and `token_dim`) and additional features such as grouped key/value (KV) sharing, optional masking, and dropout on the attention weights. **Arguments** * `query`, `key`, `value`: N‑dimensional arrays (N ≥ 4). The first three logical dimensions correspond to an embedding/head dimension, a heads dimension, and a token (sequence) dimension, but their positions are controlled by the `head_dim` and `token_dim` keyword arguments. All non-first-3 dimensions (batching dims) must match between `q`, `k`, and `v`. **Keyword arguments** * `head_dim::Int=1`: Which of the first three dimensions indexes the per-head embedding dimension. * `token_dim::Int=3`: Which of the first three dimensions indexes tokens. `head_dim` and `token_dim` must be different and must be in `1:3`. * `scale=nothing`: Scaling factor applied to the dot products. By default `scale = sqrt(size(q, head_dim))` and the implementation uses the inverse (i.e. divides by `sqrt(d_head)`). You may pass a number to override. * `mask=nothing`: Optional boolean mask broadcastable to the attention logits. When provided it selects which key positions are allowed (true = keep). * `is_causal::Union{Bool,Nothing}=nothing`: If `true` a causal mask is constructed automatically with `NNlib.make_causal_mask(...; dims=token_dim)`. `is_causal` and `mask` cannot both be specified. * `bias=nothing`: Optional additive bias applied to attention logits before softmax. * `fdrop=identity`: Function applied to the attention probabilities (e.g. a dropout function). It is applied after softmax and after mask/bias are applied. **Extended Help** **Mathematical Formulation** $$\begin{align\*} &\text{Let } q\_{d,h,l,b},\ k\_{d,h,s,b},\ v\_{e,h,s,b} \text{ denote the elements of the query, key, and value tensors,} \\ &\text{where:} \\ &\quad d: \text{embedding dimension (head\_dim)} \\ &\quad h: \text{head index} \\ &\quad l: \text{query token index} \\ &\quad s: \text{key/value token index} \\ &\quad e: \text{value embedding dimension} \\ &\quad b: \text{batch index} \\ \\ &\text{Scaling:} \\ &\quad \text{scale} = \frac{1}{\sqrt{d\_{\text{head}}}} \text{ (default, overridable)} \\ \\ &\text{Attention logits:} \\ &\quad \text{logits}*{s,l,h,b} = \text{scale} \cdot \sum*{d=1}^{d\_{\text{head}}} k\_{d,h,s,b} , q\_{d,h,l,b} + \text{bias}*{s,l,h,b} \\ &\quad \text{If mask is provided: logits}*{s,l,h,b} \rightarrow -\infty \text{ where mask is false} \\ \\ &\text{Attention weights:} \\ &\quad \alpha\_{s,l,h,b} = \text{softmax}*s(\text{logits}*{:,l,h,b}) \\ \\ &\text{Output:} \\ &\quad \text{output}*{e,l,h,b} = \sum*{s=1}^{S} v\_{e,h,s,b} , \alpha\_{s,l,h,b} \\ \end{align\*}$$ **Behavior notes** * The implementation asserts `N ≥ 4` and enforces that `q` and `k` have the same size along `head_dim`, and that `k` and `v` have the same token length (along `token_dim`). * Grouped KV: If `k`/`v` provide fewer KV groups than `q`'s number of heads, they are repeated along the non-head batching dimension so that the number of KV groups divides the number of heads. This supports e.g. shared keys for multiple heads. * Masking and bias are combined using the internal `apply_attention_mask_bias` helper which produces numerically stable logits before softmax. **Returns** * A tuple `(output, attn_weights)` where * `output` contains the attended values with the same batch/head layout as `q`/`v` (see implementation for exact ordering), and * `attn_weights` contains the attention probabilities (softmax over the key/token dimension). **Performance** * CPU: The implementation relies on `batched_matmul` (which may use LoopVectorization when available) for good CPU performance. * Reactant: Lowered to a `dot_general`-style op. It is highly recommended to use this function instead of writing out the sdpa implementation directly. source ## Batched Operations {#Batched-Operations} ```julia batched_matmul(x::AbstractMatrix, y::AbstractArray{yT,3}) where {yT} batched_matmul(x::AbstractArray{xT,3}, y::AbstractMatrix) where {xT} batched_matmul( x::AbstractArray{xT,N}, y::AbstractArray{yT,N}; lhs_contracting_dim::Int=2, rhs_contracting_dim::Int=1, lhs_batching_dims::Dims{M}=ntuple(Base.Fix2(+, 2), Val(N - 2)), rhs_batching_dims::Dims{M}=ntuple(Base.Fix2(+, 2), Val(N - 2)), ) where {xT,yT,N,M} ``` Computes the batched matrix multiplication of `x` and `y`. The following types are supported for `x` and `y`: * `AbstractMatrix` and `AbstractArray{<:Number,3}`: `x` is treated as a tensor with batch size `1`. * `AbstractArray{<:Number,3}` and `AbstractMatrix`: `y` is treated as a tensor with batch size `1`. * `AbstractArray{<:Number,N}` and `AbstractArray{<:Number,N}`: This is the general form. This case supports the keyword arguments. Size of each corresponding batch dimension must be equal (or `1`). **Keyword Arguments** * `lhs_contracting_dim`: The dimension along which `x` is contracted. Defaults to `2`. * `rhs_contracting_dim`: The dimension along which `y` is contracted. Defaults to `1`. * `lhs_batching_dims`: The batching dimensions of `x`. Defaults to the last N - 2 dimensions. * `rhs_batching_dims`: The batching dimensions of `y`. Defaults to the last N - 2 dimensions. **Output Shape** The output shape of result `z` is computed as follows: * `size(z, 1)`: size of the non-contracting dimension of `x`. * `size(z, 2)`: size of the non-contracting dimension of `y`. * `size(z)[3:end]`: size of each of the resolved batching dimensions (i.e. `max(size(x, lhs_batching_dims[j]), size(y, rhs_batching_dims[j])))`) of `x` and `y`. **Performance Considerations** * **CPU**: `x` and `y` are converted to 3D tensors and then we attempt to use a custom implementation based on LoopVectorization. If this fails, we fall back to `NNlib.batched_mul`. ::: tip Load `LoopVectorization.jl` to get faster batched matrix multiplication On CPUs loading LoopVectorization adds faster implementations of batched matrix multiplication. ::: * **GPU**: `x` and `y` are converted to 3D tensors and then we call `NNlib.batched_mul`. * **Reactant**: We directly lower this function to `stablehlo.dot_general`. source ## Bias Activation {#Bias-Activation} ```julia bias_activation(σ, x, bias) ``` Applies the activation function `σ` elementwise to the result of broadcasted addition of `x` and `bias` along the penultimate dimension. A vector `x` is treated as a matrix with a single last dimension. **Arguments** * `σ`: Activation function * `x`: Input to be transformed * `bias`: Bias to be added. Can be `nothing`. See also [`bias_activation!!`](/api/NN_Primitives/LuxLib#LuxLib.API.bias_activation!!), [`fast_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.fast_activation). source ```julia bias_activation!!(σ, x, bias) ``` Same as [`bias_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.bias_activation) but might update `x` in-place if possible. Users should not rely on `x` being mutated, it is recommended to use it like `y = bias_activation!!(σ, x, bias)`. If `x` is updated in-place, `y` aliases `x`. See also [`bias_activation`](/api/NN_Primitives/LuxLib#LuxLib.API.bias_activation), [`fast_activation!!`](/api/NN_Primitives/LuxLib#LuxLib.API.fast_activation!!). source ## Convolutional Layers {#Convolutional-Layers} ```julia fused_conv_bias_activation(σ::F, weight::AbstractArray, x::AbstractArray, b::Optional{<:AbstractVector}, cdims::ConvDims) where {F} ``` Computes `σ.(conv(x, weight, cdims) .+ b)` (`b` is not exactly broadcasted like this, rather it is reshaped and broadcasted to the penultimate dimension) with the best possible implementation available. This operation fuses operations into a single kernel if possible, and minimizes reallocations by reusing the output buffer for multiple operations. **Arguments** * `σ`: Activation function * `weight`: Weight tensor * `x`: Input tensor * `b`: Bias tensor (can be `nothing`) * `cdims`: `ConvDims` object **Notes on implementation** * For CUDA Arrays, this uses fused CUDNN kernels when the activation is `identity` or `relu`. For other activations, it tries to fuse the operations on the Julia side. * If any of the inputs, don't support setindexing (aka immutable arrays) we fallback to the generic non-mutating implementation. * Maximum memory reuse and operation fusion is guaranteed for ChainRules compatible AD backends or backends that support mutation. Backends like `Tracker` and `ReverseDiff` fallback to the generic implementation. * For Mixed-Precision Inputs on GPU, we type promote the inputs to the highest precision, with a warning. source ## Dropout {#Dropout} ```julia alpha_dropout(rng::AbstractRNG, x, p, training) alpha_dropout(rng::AbstractRNG, x, p, training, α, A, B) ``` Alpha Dropout: Dropout ensuring that the mean and variance of the output remains same as the input. For details see [Klambauer *et al.* \[8\]](/references#klambauer2017self). Use the second call signature to avoid recomputing the constants for a fixed dropout probability. **Arguments** * `rng`: Random number generator * `x`: Input Array * `p`: Probability of an element to be dropped out * `training`: Set to `Val(true)` or `True()` if running in training mode. Can be set to `nothing` to automatically determine if the function is being called within an autodiff context\` * `α`: `-1.7580993408473766`. Computed at limit x tends to infinity, `selu(x) = -λβ = α` * `A`: Scaling factor for the mean * `B`: Scaling factor for the variance **Returns** * Output Array after applying alpha dropout * Updated state for the random number generator source ```julia dropout(rng::AbstractRNG, x, p, training, invp, dims) dropout(rng::AbstractRNG, x, mask, p, training, update_mask::Union{Val, StaticBool}, invp, dims) ``` Dropout: Simple Way to prevent Neural Networks for Overfitting. For details see [Srivastava *et al.* \[9\]](/references#srivastava2014dropout). **Arguments** * `rng`: Random number generator * `x`: Input Array * `mask`: Dropout Mask. If not used then it is constructed automatically * `p`: Probability of an element to be dropped out * `training`: Set to `Val(true)` or `True()` if running in training mode. Can be set to `nothing` to automatically determine if the function is being called within an autodiff context * `update_mask`: If `Val(true)` or `True()` then the mask is generated and used. Else, the `mask` provided is directly used * `invp`: Inverse multiplied to the mask. Calculated as `invp = 1 / (1 - p)`. **Returns** * Output Array after applying dropout * Dropout Mask (if `training == false`, the returned value is meaningless) * Updated state for the random number generator source ## Fully Connected Layers {#Fully-Connected-Layers} ```julia fused_dense_bias_activation(σ::F, weight::AbstractMatrix, x::AbstractMatrix, b::Optional{<:AbstractVector}) where {F} ``` Compute `σ.(weight * x .+ b)` with the best possible implementation available. Currently this implementation attempts to minimize reallocations by reusing the output buffer for multiple operations. **Arguments** * `σ`: Activation function * `weight`: Weight matrix * `x`: Input matrix * `b`: Bias vector (can be `nothing`) **Notes on implementation** * If any of the inputs, don't support setindexing (aka immutable arrays) we fallback to the generic non-mutating implementation. * Maximum memory reuse and operation fusion is guaranteed for ChainRules compatible AD backends or backends that support mutation. Backends like `Tracker` and `ReverseDiff` fallback to the generic implementation. * For CUDA Arrays, this uses a special fused implementation via cuBLASLt. * For small CPU Arrays, we use LoopVectorization.jl. On `x86_64` we use Octavian for medium sized matrices. This is overridden if special BLAS implementations are loaded (currently `MKL`, `AppleAccelerate`, and `BLISBLAS`). ::: tip Load `Octavian.jl` Loading `Octavian.jl` enables a polyalgorithm that uses different backends based on the input sizes. ::: source ## Normalization {#Normalization} ```julia batchnorm(x, scale, bias, running_mean, running_var, training, σ=identity, momentum = 0.1f0, epsilon = eps(eltype(x)) ^ (5 // 7)) ``` Batch Normalization. For details see [Ioffe and Szegedy \[10\]](/references#ioffe2015batch). Batch Normalization computes the mean and variance for each $D\_1 \times ... \times D\_{N - 2} \times 1 \times D\_N$ input slice and normalises the input accordingly. **Arguments** * `x`: Input to be Normalized * `scale`: Scale factor ($\gamma$) (can be `nothing`) * `bias`: Bias factor ($\beta$) (can be `nothing`) * `running_mean`: Running mean (can be `nothing`) * `running_var`: Running variance (can be `nothing`) * `training`: Set to `Val(true)` or `True()` if running in training mode. Can be set to `nothing` to automatically determine if the function is being called within an autodiff context * `σ`: Activation function (default: `identity`) * `momentum`: Momentum for updating running mean and variance (default: `0.1f0`) * `epsilon`: Value added to the denominator for numerical stability (default: `eps(eltype(x)) ^ (5 / 7)`) **Returns** Normalized Array of same size as `x`. And a Named Tuple containing the updated running mean and variance. source ```julia groupnorm(x, scale, bias, groups::Int, σ::F=identity, epsilon=eps(eltype(x)) ^ (5 // 7)) ``` Group Normalization. For details see [Wu and He \[11\]](/references#wu2018group). This op is similar to batch normalization, but statistics are shared across equally-sized groups of channels and not shared across batch dimension. Thus, group normalization does not depend on the batch composition and does not require maintaining internal state for storing statistics. **Arguments** * `x`: Input to be Normalized * `scale`: Scale factor ($\gamma$) (can be `nothing`) * `bias`: Bias factor ($\beta$) (can be `nothing`) * `groups`: Number of groups * `σ`: Activation function (default: `identity`) * `epsilon`: Value added to the denominator for numerical stability (default: `eps(eltype(x)) ^ (5 / 7)`) **Returns** The normalized array is returned. source ```julia instancenorm(x, scale, bias, training, act, epsilon = eps(eltype(x)) ^ (5 // 7)) instancenorm(x, scale, bias, running_mean, running_var, training, act, momentum, epsilon = eps(eltype(x)) ^ (5 // 7)) ``` Instance Normalization. For details see [Ulyanov *et al.* \[4\]](/references#ulyanov2016instance). Instance Normalization computes the mean and variance for each $D\_1 \times ... \times D\_{N - 2} \times 1 \times 1$ input slice and normalises the input accordingly. **Arguments** * `x`: Input to be Normalized (must be atleast 3D) * `scale`: Scale factor ($\gamma$) (can be `nothing`) * `bias`: Bias factor ($\beta$) (can be `nothing`) * `running_mean`: Running mean (can be `nothing`) * `running_var`: Running variance (can be `nothing`) * `training`: Set to `Val(true)` or `True()` if running in training mode. Can be set to `nothing` to automatically determine if the function is being called within an autodiff context * `σ`: Activation function (default: `identity`) * `epsilon`: Value added to the denominator for numerical stability (default: `eps(eltype(x)) ^ (5 / 7)`) * `momentum`: Momentum for updating running mean and variance (default: `0.1f0`) **Returns** Normalized Array of same size as `x`. And a Named Tuple containing the updated running mean and variance. source ```julia layernorm(x::AbstractArray{xT, N}, scale, bias, σ = identity, dims=1:(N - 1), epsilon = eps(eltype(x)) ^ (5 / 7)) where {xT, N} ``` Layer Normalization. For details see [Ba *et al.* \[12\]](/references#ba2016layer). Given an input array $x$, this layer computes $$y = \frac{x - \mathbb{E}\[x]}{\sqrt{Var\[x] + \epsilon}} \* \gamma + \beta$$ and applies the activation function `σ` elementwise to `y`. **Arguments** * `x`: Input to be Normalized * `scale`: Scale factor ($\gamma$) (can be `nothing`) * `bias`: Bias factor ($\beta$) (can be `nothing`) * `σ`: Activation function (default: `identity`) * `dims`: Dimensions along which the mean and std of `x` is computed. If `nothing` is passed, the dims are inferred based on the dimensions of scale and bias. For example, if `x` is `N` dimensional and `scale` and `bias` are `M` dimensional, then the dims will be `1:(N - M)`. * `epsilon`: Value added to the denominator for numerical stability (default: `eps(eltype(x)) ^ (5 / 7)`) **Returns** Normalized Array of same size as `x`. source ## Helper Functions {#Helper-Functions} ```julia internal_operation_mode(xs::Tuple) internal_operation_mode(x::AbstractArray) ``` Returns the internal operation mode for the given array(s). This is useful to define custom implementations using different backends like simple Julia broadcasting, Kernel Abstractions, Loop Vectorization, etc. Currently supported modes are: * `GenericBroadcastOp`: This is the fallback for most types. For the following types this is the preferred mode: * Arrays with `fast_scalar_indexing` set to `False`. * Static Arrays * ReverseDiff Arrays * Tracker Arrays * ForwardDiff.Dual Arrays * Complex Arrays * `GPUBroadcastOp{dev}`: GPU Arrays where `dev` is obtained from `get_device_type(xs)`. This option dispatches should preferably use `KernelAbstractions` or specialized vendor dispatches. * `LoopedArrayOp`: CPU arrays that can be optimized using SIMD Loops, ideally using `LoopVectorization.jl` or `Polyester.jl`. source --- --- url: /dev/api/NN_Primitives/NNlib.md --- # NNlib {#NNlib-API} Neural Network Primitives with custom bindings for different accelerator backends in Julia. ::: tip Reexport of `NNlib` Lux doesn't re-export all of `NNlib` for now. Directly loading `NNlib` is the recommended approach for accessing these functions. ::: ## Attention {#Attention} ```julia dot_product_attention(query, key, value, [bias]; [fdrop, mask, nheads]) ``` Multihead dot product attention used in transformer architectures. The input arrays must have the first two dimensions given by the number of features and the sequence length, then an arbitrary number of batch dimensions or none. Returns the attention output array of size `(v_dim, q_len, batch_size...)` and the attention scores of size `(kv_len, q_len, nheads, batch_size...)`. See also [`dot_product_attention_scores`](/api/NN_Primitives/NNlib#NNlib.dot_product_attention_scores) if you only need the attention scores. **Arguments** * `query`: Query array of size `(qk_dim, q_len, batch_size...)`. * `key`: Key array of size `(qk_dim, kv_len, batch_size...)`. * `value`: Value array of size `(v_dim, kv_len, batch_size...)`. * `bias`: Either `nothing` or an array broadcastable to size `(kv_len, q_len, nheads, batch_size)`. It will be added to the attention scores before applying the softmax. Default `nothing`. * `fdrop`: A dropout function or layer to be applied on the attention scores right after the softmax. Default `identity` (no dropout). * `mask`: Either `nothing` or a boolean array broadcastable to size `(kv_len, q_len, nheads, batch_size)`. The mask is applied to the attention scores just before the softmax. See [`make_causal_mask`](/api/NN_Primitives/NNlib#NNlib.make_causal_mask) fore creating causal masks. Default `nothing`. * `nheads`: Number of heads to split the input arrays into. Default `1`. **Examples** ```julia q, k, v = rand(10, 20, 2), rand(10, 30, 2), rand(20, 30, 2) y, α = dot_product_attention(q, k, v) ``` source ```julia dot_product_attention_scores(query, key, [bias]; [fdrop, mask]) ``` Return the attention scores for the [`dot_product_attention`](/api/NN_Primitives/NNlib#NNlib.dot_product_attention). Input arrays must have dimensions `(num_features ÷ nheads, nheads, sequence_length, batch_size)`. See [`dot_product_attention`](/api/NN_Primitives/NNlib#NNlib.dot_product_attention) for more details. source ```julia make_causal_mask(x, dims=2) ``` Return a boolean square matrix `m` of the same type as `x` and of side `size(x, dims)`. Its elements are set such that `m[i, j] == i ≤ j`. Can be used to mask the attention scores in [`dot_product_attention`](/api/NN_Primitives/NNlib#NNlib.dot_product_attention). source ## Softmax {#Softmax} ```julia softmax(x; dims = 1) ``` [Softmax](https://en.wikipedia.org/wiki/Softmax_function) turns input array `x` into probability distributions that sum to 1 along the dimensions specified by `dims`. It is semantically equivalent to the following: ```julia softmax(x; dims = 1) = exp.(x) ./ sum(exp.(x), dims = dims) ``` with additional manipulations enhancing numerical stability. For a matrix input `x` it will by default (`dims = 1`) treat it as a batch of vectors, with each column independent. Keyword `dims = 2` will instead treat rows independently, and so on. See also [`logsoftmax`](/api/NN_Primitives/NNlib#NNlib.logsoftmax). **Examples** ```julia julia> softmax([1, 2, 3]) 3-element Vector{Float64}: 0.09003057317038046 0.24472847105479764 0.6652409557748218 julia> softmax([1 2 3; 2 2 2]) # dims=1 2×3 Matrix{Float64}: 0.268941 0.5 0.731059 0.731059 0.5 0.268941 julia> softmax([1 2 3; 2 2 2]; dims=2) 2×3 Matrix{Float64}: 0.0900306 0.244728 0.665241 0.333333 0.333333 0.333333 ``` Note that, when used with Flux.jl, `softmax` must not be passed to layers like `Dense` which accept an activation function. The activation is broadcasted over the result, thus applies to individual numbers. But `softmax` always needs to see the whole column. ```julia julia> using Flux julia> x = randn(Float32, 4, 4, 3, 13); julia> model = Chain(Conv((4, 4), 3 => 8, tanh), Flux.flatten, Dense(8 => 7), softmax); julia> model(x) |> size (7, 13) julia> Dense(4 => 7, softmax)(x) ERROR: `softmax(x)` called with a number, but it expects an array. ``` source ```julia logsoftmax(x; dims = 1) ``` Computes the log of softmax in a more numerically stable way than directly taking `log.(softmax(xs))`. Commonly used in computing cross entropy loss. It is semantically equivalent to the following: ```julia logsoftmax(x; dims = 1) = x .- log.(sum(exp.(x), dims = dims)) ``` See also [`softmax`](/api/NN_Primitives/NNlib#NNlib.softmax). source ## Pooling {#Pooling} ```julia PoolDims(x_size::NTuple{M}, k::Union{NTuple{L, Int}, Int}; stride=k, padding=0, dilation=1) where {M, L} ``` Dimensions for a "pooling" operation that can have an arbitrary input size, kernel size, stride, dilation, and channel count. Used to dispatch onto efficient implementations at compile-time. source ```julia maxpool(x, k::NTuple{N, Integer}; pad=0, stride=k) ``` Perform max pool operation with window size `k` on input tensor `x`. Arguments: * `x` and `k`: Expects `ndim(x) ∈ 3:5`, and always `length(k) == ndim(x) - 2` * `pad`: See [`pad_zeros`](/api/NN_Primitives/NNlib#NNlib.pad_zeros) for details. * `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`. source ```julia meanpool(x, k::NTuple{N, Integer}; pad=0, stride=k) ``` Perform mean pool operation with window size `k` on input tensor `x`. Arguments: * `x` and `k`: Expects `ndim(x) ∈ 3:5`, and always `length(k) == ndim(x) - 2` * `pad`: See [`pad_zeros`](/api/NN_Primitives/NNlib#NNlib.pad_zeros) for details. * `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`. source ```julia lpnormpool(x, p::Real, k::NTuple{N, Integer}; pad=0, stride=k) ``` Perform Lp pool operation with value of the Lp norm `p` and window size `k` on input tensor `x`, also known as LPPool in pytorch. This pooling operator from [Learned-Norm Pooling for Deep Feedforward and Recurrent Neural Networks](https://arxiv.org/abs/1311.1780). Arguments: * `x` and `k`: Expects `ndim(x) ∈ 3:5`, and always `length(k) == ndim(x) - 2` * `p` is restricted to `0 < p < Inf`. * `pad`: See [`pad_zeros`](/api/NN_Primitives/NNlib#NNlib.pad_zeros) for details. * `stride`: Either a tuple with the same length as `k`, or one integer for all directions. Default is `k`. For all elements `x` in a size `k` window, lpnormpool computes `(∑ᵢ xᵢ^p)^(1 / p)` as an element of the output. Thus `lpnormpool(x, 1, k) ./ prod(k) ≈ meanpool(x, k)` and `lpnormpool(x, 2, k).^2 ./ prod(k) ≈ meanpool(x.^2, k)`. source ## Padding {#Padding} ```julia pad_reflect(x, pad::Tuple; [dims]) pad_reflect(x, pad::Int; [dims]) ``` Pad the array `x` reflecting its values across the border. `pad` can a tuple of integers `(l1, r1, ..., ln, rn)` of some length `2n` that specifies the left and right padding size for each of the dimensions in `dims`. If `dims` is not given, it defaults to the first `n` dimensions. If `pad` is an integer, it is applied on both sides on every dimension in `dims`. In this case, `dims` defaults to the first `ndims(x)-2` dimensions (i.e. excludes the channel and batch dimension). See also [`pad_repeat`](/api/NN_Primitives/NNlib#NNlib.pad_repeat), [`pad_symmetric`](/api/NN_Primitives/NNlib#NNlib.pad_symmetric), [`pad_circular`](/api/NN_Primitives/NNlib#NNlib.pad_circular), and [`pad_constant`](/api/NN_Primitives/NNlib#NNlib.pad_constant). ```julia julia> r = reshape(1:9, 3, 3) 3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64: 1 4 7 2 5 8 3 6 9 julia> pad_reflect(r, (1,2,1,2)) 6×6 Matrix{Int64}: 5 2 5 8 5 2 4 1 4 7 4 1 5 2 5 8 5 2 6 3 6 9 6 3 5 2 5 8 5 2 4 1 4 7 4 1 ``` source ```julia pad_symmetric(x, pad::Tuple; [dims]) pad_symmetric(x, pad::Int; [dims]) ``` Pad the array `x` reflecting its values symmetrically across the border, i.e. the border values of `x` are present in the padding values, in contrast to [`pad_reflect`](/api/NN_Primitives/NNlib#NNlib.pad_reflect). `pad` can a tuple of integers `(l1, r1, ..., ln, rn)` of some length `2n` that specifies the left and right padding size for each of the dimensions in `dims`. If `dims` is not given, it defaults to the first `n` dimensions. If `pad` is an integer, it is applied on both sides on every dimension in `dims`. In this case, `dims` defaults to the first `ndims(x)-2` dimensions (i.e. excludes the channel and batch dimension). See also [`pad_repeat`](/api/NN_Primitives/NNlib#NNlib.pad_repeat), [`pad_reflect`](/api/NN_Primitives/NNlib#NNlib.pad_reflect), [`pad_circular`](/api/NN_Primitives/NNlib#NNlib.pad_circular), and [`pad_constant`](/api/NN_Primitives/NNlib#NNlib.pad_constant). ```julia julia> r = reshape(1:9, 3, 3) 3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64: 1 4 7 2 5 8 3 6 9 julia> pad_symmetric(r, (1,2,1,2)) 6×6 Matrix{Int64}: 1 1 4 7 7 4 1 1 4 7 7 4 2 2 5 8 8 5 3 3 6 9 9 6 3 3 6 9 9 6 2 2 5 8 8 5 ``` source ```julia pad_circular(x, pad::Tuple; [dims]) pad_circular(x, pad::Int; [dims]) ``` Pad the array `x` "circularly" across the border by wrapping around values from the opposite side of `x`. `pad` can a tuple of integers `(l1, r1, ..., ln, rn)` of some length `2n` that specifies the left and right padding size for each of the dimensions in `dims`. If `dims` is not given, it defaults to the first `n` dimensions. If `pad` is an integer, it is applied on both sides on every dimension in `dims`. In this case, `dims` defaults to the first `ndims(x)-2` dimensions (i.e. excludes the channel and batch dimension). The pad length on either side in any dimension must not exceed the size of `x` in that dimension, i.e. `pad_circular` is not able to create abitrary sized tilings of `x`. See also [`pad_repeat`](/api/NN_Primitives/NNlib#NNlib.pad_repeat), [`pad_reflect`](/api/NN_Primitives/NNlib#NNlib.pad_reflect), [`pad_symmetric`](/api/NN_Primitives/NNlib#NNlib.pad_symmetric), and [`pad_constant`](/api/NN_Primitives/NNlib#NNlib.pad_constant). ```julia julia> r = reshape(1:9, 3, 3) 3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64: 1 4 7 2 5 8 3 6 9 julia> pad_circular(r, (1,2,1,2)) 6×6 Matrix{Int64}: 9 3 6 9 3 6 7 1 4 7 1 4 8 2 5 8 2 5 9 3 6 9 3 6 7 1 4 7 1 4 8 2 5 8 2 5 ``` source ```julia pad_repeat(x, pad::Tuple; [dims]) pad_repeat(x, pad::Int; [dims]) ``` Pad the array `x` repeating the values on the border. `pad` can a tuple of integers `(l1, r1, ..., ln, rn)` of some length `2n` that specifies the left and right padding size for each of the dimensions in `dims`. If `dims` is not given, it defaults to the first `n` dimensions. If `pad` is an integer, it is applied on both sides on every dimension in `dims`. In this case, `dims` defaults to the first `ndims(x)-2` dimensions (i.e. excludes the channel and batch dimension). See also [`pad_reflect`](/api/NN_Primitives/NNlib#NNlib.pad_reflect), [`pad_symmetric`](/api/NN_Primitives/NNlib#NNlib.pad_symmetric), [`pad_circular`](/api/NN_Primitives/NNlib#NNlib.pad_circular), and [`pad_constant`](/api/NN_Primitives/NNlib#NNlib.pad_constant). ```julia julia> r = reshape(1:9, 3, 3) 3×3 reshape(::UnitRange{Int64}, 3, 3) with eltype Int64: 1 4 7 2 5 8 3 6 9 julia> pad_repeat(r, (1,2,3,4)) 6×10 Matrix{Int64}: 1 1 1 1 4 7 7 7 7 7 1 1 1 1 4 7 7 7 7 7 2 2 2 2 5 8 8 8 8 8 3 3 3 3 6 9 9 9 9 9 3 3 3 3 6 9 9 9 9 9 3 3 3 3 6 9 9 9 9 9 ``` source ```julia pad_constant(x, pad::Tuple, val = 0; [dims = :]) pad_constant(x, pad::Int, val = 0; [dims = :]) ``` Pad the array `x` with the constant value `val`. `pad` can be a tuple of integers. If it is of some length `2 * length(dims)` that specifies the left and right padding size for each of the dimensions in `dims` as `(l1, r1, ..., ln, rn)`. If supplied with a tuple of length `length(dims)` instead, it applies symmetric padding. If `dims` is not given, it defaults to all dimensions. For integer `pad` input, it is applied on both sides on every dimension in `dims`. See also [`pad_zeros`](/api/NN_Primitives/NNlib#NNlib.pad_zeros), [`pad_repeat`](/api/NN_Primitives/NNlib#NNlib.pad_repeat), [`pad_reflect`](/api/NN_Primitives/NNlib#NNlib.pad_reflect), [`pad_symmetric`](/api/NN_Primitives/NNlib#NNlib.pad_symmetric), and [`pad_circular`](/api/NN_Primitives/NNlib#NNlib.pad_circular). ```julia julia> r = reshape(1:4, 2, 2) 2×2 reshape(::UnitRange{Int64}, 2, 2) with eltype Int64: 1 3 2 4 julia> pad_constant(r, (1, 2, 3, 4), 8) 5×9 Matrix{Int64}: 8 8 8 8 8 8 8 8 8 8 8 8 1 3 8 8 8 8 8 8 8 2 4 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 8 julia> pad_constant(r, 1, 8) 4×4 Matrix{Int64}: 8 8 8 8 8 1 3 8 8 2 4 8 8 8 8 8 julia> r = reshape(1:27, 3, 3, 3) 3×3×3 reshape(::UnitRange{Int64}, 3, 3, 3) with eltype Int64: [:, :, 1] = 1 4 7 2 5 8 3 6 9 [:, :, 2] = 10 13 16 11 14 17 12 15 18 [:, :, 3] = 19 22 25 20 23 26 21 24 27 julia> pad_constant(r, (2,1), dims = 1) # assymetric padding 6×3×3 Array{Int64, 3}: [:, :, 1] = 0 0 0 0 0 0 1 4 7 2 5 8 3 6 9 0 0 0 [:, :, 2] = 0 0 0 0 0 0 10 13 16 11 14 17 12 15 18 0 0 0 [:, :, 3] = 0 0 0 0 0 0 19 22 25 20 23 26 21 24 27 0 0 0 julia> pad_constant(r, (2,1, 3), dims = (1,2)) # padding must always be either the same length as dims, or double it ERROR: ArgumentError: Could not parse padding (2, 1, 3) and dims (1, 2) Stacktrace: [...] ``` source ```julia pad_zeros(x, pad::Tuple; [dims]) pad_zeros(x, pad::Int; [dims]) ``` Pad the array `x` with zeros. Equivalent to [`pad_constant`](/api/NN_Primitives/NNlib#NNlib.pad_constant) with the constant equal to 0. source ## Convolution {#Convolution} ```julia conv(x, w; stride = 1, pad = 0, dilation = 1, flipped = false, groups = 1) ``` Apply convolution filter `w` to input `x`. `x` and `w` are 3d/4d/5d tensors in 1d/2d/3d convolutions respectively. `x` and `w` may have real or complex element types. source ```julia ConvDims ``` Type system-level information about convolution dimensions. Critical for things like `im2col!()` to generate efficient code, and helpful to reduce the number of kwargs getting passed around. source ```julia depthwiseconv(x, w; stride=1, pad=0, dilation=1, flipped=false) ``` Depthwise convolution operation with filter `w` on input `x`. `x` and `w` are 3d/4d/5d tensors in 1d/2d/3d convolutions respectively. source ```julia DepthwiseConvDims ``` Concrete subclass of `ConvDims` for a depthwise convolution. Differs primarily due to characterization by `C_in`, `C_mult`, rather than `C_in`, `C_out`. Useful to be separate from DenseConvDims primarily for channel calculation differences. source ```julia DenseConvDims ``` Concrete subclass of `ConvDims` for a normal, dense, conv2d/conv3d. source ```julia unfold(x, kernel_size; stride = 1, pad = 0, dilation = 0, flipped = true) ``` Places sliding windows of x into a container tensor of size `(num_windows, window_size, batchsize)`. The window size is determined by the `prod(spatial dims of kernel)*input_channels`. The number of sliding windows will match those of convolution (`conv`) with the same kernel\_size and arguments. Note that by default `conv` flips the spatial dimensions of its kernel (default `flipped=false`), whereas `unfold` does not (default `flipped=true`). Uses `NNlib.im2col!` as backend. See also [`fold`](/api/NN_Primitives/NNlib#NNlib.fold), the adjoint/transpose operator and a potential inverse of `unfold`. **Example** The below example demonstrates that `unfold` uses the same sliding windows as `conv`. In general [`batched_mul`](/api/NN_Primitives/NNlib#NNlib.batched_mul) + `unfold` should not be used to achieve convolution. ```julia julia> x = reshape([100 2 3 40 5 6 700], 7, 1, 1); # 1D data, 1 channel, batch of 1 julia> w = reshape([1 0 -1], 3, 1, 1); # 1D conv kernel of length 3 julia> kws = (pad=1, stride=2, flipped=true); # use same args for conv and unfold julia> z = NNlib.unfold(x, size(w); kws...) 4×3×1 Array{Int64, 3}: [:, :, 1] = 0 100 2 2 3 40 40 5 6 6 700 0 julia> y1 = conv(x, w; kws...) 4×1×1 Array{Int64, 3}: [:, :, 1] = -2 -38 34 6 julia> y2 = z ⊠ w # ⊠ (\boxtimes) is NNlib.batched_mul 4×1×1 Array{Int64, 3}: [:, :, 1] = -2 -38 34 6 ``` source ```julia fold(y, output_size, kernel_size; stride = 1, pad = 0, dilation = 0, flipped = true) ``` The adjoint/transpose operator of `unfold`. It accumulates sliding windows from the output of `unfold` into a container tensor of size `output_size`. An inverse to `unfold` may be obtained (in some cases) by using `fold` and accounting for scaling issues with a divisor (see example). Uses `NNlib.col2im!` as backend. See also [`unfold`](/api/NN_Primitives/NNlib#NNlib.unfold). **Example** ```julia julia> x = reshape([100 2 3 40 5 6 700], 7, 1, 1); # 1D data, 1 channel, batch of 1 julia> y = NNlib.unfold(x, (3,1,1)) # sliding window of size 3 5×3×1 Array{Int64, 3}: [:, :, 1] = 100 2 3 2 3 40 3 40 5 40 5 6 5 6 700 julia> z = NNlib.fold(y, size(x), (3,1,1)) # sum of contributions in y. 100 appears once, 40 three times 7×1×1 Array{Int64, 3}: [:, :, 1] = 100 4 9 120 15 12 700 julia> divisor = NNlib.fold(NNlib.unfold(ones(size(x)...), (3,1,1)), size(x), (3,1,1)) 7×1×1 Array{Float64, 3}: [:, :, 1] = 1.0 2.0 3.0 3.0 3.0 2.0 1.0 julia> z ./ divisor 7×1×1 Array{Float64, 3}: [:, :, 1] = 100.0 2.0 3.0 40.0 5.0 6.0 700.0 ``` In general, an inverse to `unfold` does not exist if `divisor` contains zeros. source ## Upsampling {#Upsampling} ```julia upsample_nearest(x, scale::NTuple{S,Int}) upsample_nearest(x; size::NTuple{S,Int}) ``` Upsamples the array `x` by integer multiples along the first `S` dimensions. Subsequent dimensions of `x` are not altered. Either the `scale` factors or the final output `size` can be specified. See also [`upsample_bilinear`](/api/NN_Primitives/NNlib#NNlib.upsample_bilinear), for two dimensions of an `N=4` array. **Example** ```julia julia> upsample_nearest([1 2 3; 4 5 6], (2, 3)) 4×9 Matrix{Int64}: 1 1 1 2 2 2 3 3 3 1 1 1 2 2 2 3 3 3 4 4 4 5 5 5 6 6 6 4 4 4 5 5 5 6 6 6 julia> ans == upsample_nearest([1 2 3; 4 5 6]; size=(4, 9)) # equivalent true julia> upsample_nearest([1 2 3; 4 5 6], (2,)) 4×3 Matrix{Int64}: 1 2 3 1 2 3 4 5 6 4 5 6 julia> ans == upsample_nearest([1 2 3; 4 5 6], size=(4,)) true ``` source ```julia ∇upsample_nearest(Δ::AbstractArray{T,3}, scales::NTuple{S, <:Integer}) where T ``` **Arguments** * `Δ`: Incoming gradient array, backpropagated from downstream layers * `scales`: scales by which the image was upsampled in the first place **Outputs** * `dx`: Downsampled version of `Δ` source ```julia upsample_linear(x::AbstractArray{T,3}, scale::Real; align_corners::Bool = true) upsample_linear(x::AbstractArray{T,3}; size::Integer, align_corners::Bool = true) ``` Upsamples the first dimension of the array `x` by the upsample provided `scale`, using linear interpolation. As an alternative to using `scale`, the resulting array `size` can be directly specified with a keyword argument. The size of the output is equal to `(scale*S1, S2, S3)`, where `S1, S2, S3 = size(x)`. source ```julia ∇upsample_linear(Δ::AbstractArray{T,3}; size::Integer, align_corners::Bool = true) where T ``` **Arguments** * `Δ`: Incoming gradient array, backpropagated from downstream layers * `size`: Size of the image upsampled in the first place **Outputs** * `dx`: Downsampled version of `Δ` source ```julia upsample_bilinear(x::AbstractArray{T,4}, scale::NTuple{2,Real}; align_corners::Bool = true) upsample_bilinear(x::AbstractArray{T,4}; size::NTuple{2,Integer}, align_corners::Bool = true) ``` Upsamples the first 2 dimensions of the array `x` by the upsample factors stored in `scale`, using bilinear interpolation. As an alternative to using `scale`, the resulting image `size` can be directly specified with a keyword argument. The size of the output is equal to `(scale[1]*S1, scale[2]*S2, S3, S4)`, where `S1, S2, S3, S4 = size(x)`. **Examples** ```julia julia> x = reshape(Float32[1 2 3; 4 5 6], (2,3,1,1)) 2×3×1×1 Array{Float32, 4}: [:, :, 1, 1] = 1.0 2.0 3.0 4.0 5.0 6.0 julia> upsample_bilinear(x, (2, 3)) 4×9×1×1 Array{Float32, 4}: [:, :, 1, 1] = 1.0 1.25 1.5 1.75 2.0 2.25 2.5 2.75 3.0 2.0 2.25 2.5 2.75 3.0 3.25 3.5 3.75 4.0 3.0 3.25 3.5 3.75 4.0 4.25 4.5 4.75 5.0 4.0 4.25 4.5 4.75 5.0 5.25 5.5 5.75 6.0 julia> ans == upsample_bilinear(x; size=(4, 9)) # specify ouput size instead true julia> upsample_bilinear(x, (2.5, 3.5)) # non-integer scaling factors are allowed 5×10×1×1 Array{Float32, 4}: [:, :, 1, 1] = 1.0 1.22222 1.44444 1.66667 1.88889 … 2.33333 2.55556 2.77778 3.0 1.75 1.97222 2.19444 2.41667 2.63889 3.08333 3.30556 3.52778 3.75 2.5 2.72222 2.94444 3.16667 3.38889 3.83333 4.05556 4.27778 4.5 3.25 3.47222 3.69444 3.91667 4.13889 4.58333 4.80556 5.02778 5.25 4.0 4.22222 4.44444 4.66667 4.88889 5.33333 5.55556 5.77778 6.0 ``` source ```julia ∇upsample_bilinear(Δ::AbstractArray{T,4}; size::NTuple{2,Integer}, align_corners::Bool = true) where T ``` **Arguments** * `Δ`: Incoming gradient array, backpropagated from downstream layers * `size`: Lateral (W,H) size of the image upsampled in the first place **Outputs** * `dx`: Downsampled version of `Δ` source ```julia upsample_trilinear(x::AbstractArray{T,5}, scale::NTuple{3,Real}; align_corners::Bool = true) upsample_trilinear(x::AbstractArray{T,5}; size::NTuple{3,Integer}, align_corners::Bool = true) ``` Upsamples the first 3 dimensions of the array `x` by the upsample factors stored in `scale`, using trilinear interpolation. As an alternative to using `scale`, the resulting image `size` can be directly specified with a keyword argument. The size of the output is equal to `(scale[1]*S1, scale[2]*S2, scale[3]*S3, S4, S5)`, where `S1, S2, S3, S4, S5 = size(x)`. **Examples** ```julia upsample_trilinear(x, (2, 3, 4)) upsample_trilinear(x; size=(4, 9, 11)) # specify ouput size instead upsample_trilinear(x, (2.5, 3.5, pi)) # non-integer scaling factors are allowed ``` source ```julia ∇upsample_trilinear(Δ::AbstractArray{T,5}; size::NTuple{3,Integer}, align_corners::Bool = true) where T ``` **Arguments** * `Δ`: Incoming gradient array, backpropagated from downstream layers * `size`: Lateral size & depth (W,H,D) of the image upsampled in the first place **Outputs** * `dx`: Downsampled version of `Δ` source ```julia pixel_shuffle(x, r::Integer) ``` Pixel shuffling operation, upscaling by a factor `r`. For 4-arrays representing `N` images, the operation converts input `size(x) == (W, H, r^2*C, N)` to output of size `(r*W, r*H, C, N)`. For `D`-dimensional data, it expects `ndims(x) == D+2` with channel and batch dimensions, and divides the number of channels by `r^D`. Used in super-resolution networks to upsample towards high resolution features. Reference: Shi et. al., "Real-Time Single Image and Video Super-Resolution ...", CVPR 2016, https://arxiv.org/abs/1609.05158 **Examples** ```julia julia> x = [10i + j + channel/10 for i in 1:2, j in 1:3, channel in 1:4, batch in 1:1] 2×3×4×1 Array{Float64, 4}: [:, :, 1, 1] = 11.1 12.1 13.1 21.1 22.1 23.1 [:, :, 2, 1] = 11.2 12.2 13.2 21.2 22.2 23.2 [:, :, 3, 1] = 11.3 12.3 13.3 21.3 22.3 23.3 [:, :, 4, 1] = 11.4 12.4 13.4 21.4 22.4 23.4 julia> pixel_shuffle(x, 2) # 4 channels used up as 2x upscaling of image dimensions 4×6×1×1 Array{Float64, 4}: [:, :, 1, 1] = 11.1 11.3 12.1 12.3 13.1 13.3 11.2 11.4 12.2 12.4 13.2 13.4 21.1 21.3 22.1 22.3 23.1 23.3 21.2 21.4 22.2 22.4 23.2 23.4 julia> y = [i + channel/10 for i in 1:3, channel in 1:6, batch in 1:1] 3×6×1 Array{Float64, 3}: [:, :, 1] = 1.1 1.2 1.3 1.4 1.5 1.6 2.1 2.2 2.3 2.4 2.5 2.6 3.1 3.2 3.3 3.4 3.5 3.6 julia> pixel_shuffle(y, 2) # 1D image, with 6 channels reduced to 3 6×3×1 Array{Float64, 3}: [:, :, 1] = 1.1 1.3 1.5 1.2 1.4 1.6 2.1 2.3 2.5 2.2 2.4 2.6 3.1 3.3 3.5 3.2 3.4 3.6 ``` source ## Rotation {#Rotation} Rotate images in the first two dimensions of an array. ```julia imrotate(arr::AbstractArray{T, 4}, θ; method=:bilinear, rotation_center=size(arr) .÷ 2 .+ 1) ``` Rotates an array in the first two dimensions around the center pixel `rotation_center`. The default value of `rotation_center` is defined such that there is a integer center pixel for even and odd sized arrays which it is rotated around. For an even sized array of size `(4,4)` this would be `(3,3)`, for an odd array of size `(3,3)` this would be `(2,2)` However, `rotation_center` can be also non-integer numbers if specified. The angle `θ` is interpreted in radians. The adjoint is defined with ChainRulesCore.jl. This method also runs with CUDA (and in principle all KernelAbstractions.jl supported backends). **Keywords** * `method=:bilinear` for bilinear interpolation or `method=:nearest` for nearest neighbour * `rotation_center=size(arr) .÷ 2 .+ 1` means there is a real center pixel around it is rotated. **Examples** ```julia julia> arr = zeros((4,4,1,1)); arr[2,2,1,1] = 1; julia> arr 4×4×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 julia> NNlib.imrotate(arr, deg2rad(90)) # rotation around (3,3) 4×4×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 0.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 julia> NNlib.imrotate(arr, deg2rad(90), rotation_center=(2,2)) 4×4×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 julia> arr = zeros((3,3,1,1)); arr[1,2,1,1] = 1 1 julia> arr 3×3×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 julia> NNlib.imrotate(arr, deg2rad(45)) 3×3×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 0.207107 0.0 0.0 0.0 0.207107 0.0 0.0 0.0 julia> NNlib.imrotate(arr, deg2rad(45), method=:nearest) 3×3×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 0.0 ``` source ```julia ∇imrotate(dy, arr::AbstractArray{T, 4}, θ; method=:bilinear, rotation_center=size(arr) .÷ 2 .+ 1) ``` Adjoint for `imrotate`. Gradient only with respect to `arr` and not `θ`. **Arguments** * `dy`: input gradient * `arr`: Input from primal computation * `θ`: rotation angle in radians * `method=:bilinear` or `method=:nearest` * `rotation_center=size(arr) .÷ 2 .+ 1` rotates around a real center pixel for even and odd sized arrays source ## Batched Operations {#Batched-Operations} ```julia batched_mul(A, B) -> C A ⊠ B # \boxtimes ``` Batched matrix multiplication. Result has `C[:,:,k...] == A[:,:,k...] * B[:,:,k...]` where `k...` represent any indices in the last dimensions. If `ndims(A) == ndims(B) == 3` and `size(B,3) == 1` then instead `C[:,:,k] == A[:,:,k] * B[:,:,1]`, and similarly for `A`. To transpose each matrix, apply `batched_transpose` to the array, or `batched_adjoint` for conjugate-transpose: ```julia julia> A, B = randn(2,5,17), randn(5,9,17); julia> A ⊠ B |> size (2, 9, 17) julia> batched_adjoint(A) |> size (5, 2, 17) julia> batched_mul(A, batched_adjoint(randn(9,5,17))) |> size (2, 9, 17) julia> A ⊠ randn(5,9,1) |> size (2, 9, 17) julia> batched_transpose(A) == PermutedDimsArray(A, (2,1,3)) true ``` The equivalent `PermutedDimsArray` may be used in place of `batched_transpose`. Other permutations are also handled by BLAS, provided that the batch index `k` is not the first dimension of the underlying array. Thus `PermutedDimsArray(::Array, (1,3,2))` and `PermutedDimsArray(::Array, (3,1,2))` are fine. However, `A = PermutedDimsArray(::Array, (3,2,1))` is not acceptable to BLAS, since the batch dimension is the contiguous one: `stride(A,3) == 1`. This will be copied, as doing so is faster than `batched_mul_generic!`. Both this `copy` and `batched_mul_generic!` produce `@debug` messages, and setting for instance `ENV["JULIA_DEBUG"] = NNlib` will display them. source ```julia batched_mul(A::Array{T,3}, B::Matrix) batched_mul(A::Matrix, B::Array{T,3}) A ⊠ B ``` This is always matrix-matrix multiplication, but either `A` or `B` may lack a batch index. * When `B` is a matrix, result has `C[:,:,k] == A[:,:,k] * B[:,:]` for all `k`. * When `A` is a matrix, then `C[:,:,k] == A[:,:] * B[:,:,k]`. This can also be done by reshaping and calling `*`, for instance `A ⊡ B` using TensorCore.jl, but is implemented here using `batched_gemm` instead of `gemm`. ```julia julia> randn(16,8,32) ⊠ randn(8,4) |> size (16, 4, 32) julia> randn(16,8,32) ⊠ randn(8,4,1) |> size # equivalent (16, 4, 32) julia> randn(16,8) ⊠ randn(8,4,32) |> size (16, 4, 32) ``` See also `batched_vec` to regard `B` as a batch of vectors, `A[:,:,k] * B[:,k]`. source ```julia batched_mul!(C, A, B) -> C batched_mul!(C, A, B, α=1, β=0) ``` In-place batched matrix multiplication, equivalent to `mul!(C[:,:,k], A[:,:,k], B[:,:,k], α, β)` for all `k`. If `size(B,3) == 1` then every batch uses `B[:,:,1]` instead. This will call `batched_gemm!` whenever possible. For real arrays this means that, for `X ∈ [A,B,C]`, either `stride(X,1)==1` or `stride(X,2)==1`, the latter may be caused by `batched_transpose` or by for instance `PermutedDimsArray(::Array, (3,1,2))`. Unlike `batched_mul` this will never make a copy. For complex arrays, the wrapper made by `batched_adjoint` must be outermost to be seen. In this case the strided accepted by BLAS are more restricted, if `stride(C,1)==1` then only `stride(AorB::BatchedAdjoint,2) == 1` is accepted. source ```julia batched_transpose(A::AbstractArray{T,3}) batched_adjoint(A) ``` Equivalent to applying `transpose` or `adjoint` to each matrix `A[:,:,k]`. These exist to control how `batched_mul` behaves, as it operates on such matrix slices of an array with `ndims(A)==3`. `PermutedDimsArray(A, (2,1,3))` is equivalent to `batched_transpose(A)`, and is also understood by `batched_mul` (and more widely supported elsewhere). ```julia BatchedTranspose{T, S} <: AbstractBatchedMatrix{T, 3} BatchedAdjoint{T, S} ``` Lazy wrappers analogous to `Transpose` and `Adjoint`, returned by `batched_transpose` etc. source ```julia batched_transpose(A::AbstractArray{T,3}) batched_adjoint(A) ``` Equivalent to applying `transpose` or `adjoint` to each matrix `A[:,:,k]`. These exist to control how `batched_mul` behaves, as it operates on such matrix slices of an array with `ndims(A)==3`. `PermutedDimsArray(A, (2,1,3))` is equivalent to `batched_transpose(A)`, and is also understood by `batched_mul` (and more widely supported elsewhere). ```julia BatchedTranspose{T, S} <: AbstractBatchedMatrix{T, 3} BatchedAdjoint{T, S} ``` Lazy wrappers analogous to `Transpose` and `Adjoint`, returned by `batched_transpose` etc. source ```julia batched_vec(A::AbstractArray{T,3}, B::AbstractMatrix) batched_vec(A::AbstractArray{T,3}, b::AbstractVector) batched_vec(A::AbstractArray, B::AbstractArray) ``` Batched matrix-vector multiplication. For the 3D case: the result has `C[:,:,k] == A[:,:,k] * B[:,k]` for all `k`, or else `C[:,:,k] == A[:,:,k] * b` for `b::Vector`. For the general N-D case where `ndims(A) == ndims(B) + 1`: the result has `C[:,k...] == A[:,:,k...] * B[:,k...]` for all batch indices `k...`. The batch dimensions must match: `size(A)[3:end] == size(B)[2:end]`. With the same argument types, `batched_mul(A, B)` would regard `B` as a fixed matrix, not a batch of vectors. Both reshape and then call `batched_mul(::Array{T,3}, ::Array{T,3})`. ```julia julia> A, B, b = randn(16,8,32), randn(8,32), randn(8); julia> batched_vec(A,B) |> size (16, 32) julia> batched_vec(A,b) |> size (16, 32) julia> A4d, B3d = randn(16,8,10,32), randn(8,10,32); # 4D and 3D arrays julia> batched_vec(A4d, B3d) |> size (16, 10, 32) ``` source ## Gather and Scatter {#Gather-and-Scatter} ```julia NNlib.gather(src, idx) -> dst ``` Reverse operation of [`scatter`](/api/NN_Primitives/NNlib#NNlib.scatter). Gathers data from source `src` and writes it in a destination `dst` according to the index array `idx`. For each `k` in `CartesianIndices(idx)`, assign values to `dst` according to ```julia dst[:, ... , k] .= src[:, ... , idx[k]...] ``` Notice that if `idx` is a vector containing integers and `src` is a matrix, previous expression simplifies to ```julia dst[:, k] .= src[:, idx[k]] ``` and `k` will run over `1:length(idx)`. The elements of `idx` can be integers or integer tuples and may be repeated. A single `src` column can end up being copied into zero, one, or multiple `dst` columns. See [`gather!`](/api/NN_Primitives/NNlib#NNlib.gather!) for an in-place version. **Examples** ```julia julia> NNlib.gather([1,20,300,4000], [2,4,2]) 3-element Vector{Int64}: 20 4000 20 julia> NNlib.gather([1 2 3; 4 5 6], [1,3,1,3,1]) 2×5 Matrix{Int64}: 1 3 1 3 1 4 6 4 6 4 ``` source ```julia gather(src, IJK...) ``` Convert the tuple of integer vectors `IJK` to a tuple of `CartesianIndex` and call `gather` on it: `gather(src, CartesianIndex.(IJK...))`. **Examples** ```julia julia> src = reshape([1:15;], 3, 5) 3×5 Matrix{Int64}: 1 4 7 10 13 2 5 8 11 14 3 6 9 12 15 julia> NNlib.gather(src, [1, 2], [2, 4]) 2-element Vector{Int64}: 4 11 ``` source ```julia NNlib.gather!(dst, src, idx) ``` Reverse operation of [`scatter!`](/api/NN_Primitives/NNlib#NNlib.scatter!). Gathers data from source `src` and writes it in destination `dst` according to the index array `idx`. For each `k` in `CartesianIndices(idx)`, assign values to `dst` according to ```julia dst[:, ... , k] .= src[:, ... , idx[k]...] ``` Notice that if `idx` is a vector containing integers, and both `dst` and `src` are matrices, previous expression simplifies to ```julia dst[:, k] .= src[:, idx[k]] ``` and `k` will run over `1:length(idx)`. The elements of `idx` can be integers or integer tuples and may be repeated. A single `src` column can end up being copied into zero, one, or multiple `dst` columns. See [`gather`](/api/NN_Primitives/NNlib#NNlib.gather) for an allocating version. source ```julia NNlib.scatter(op, src, idx; [init, dstsize]) ``` Scatter operation allocating a destination array `dst` and calling `scatter!(op, dst, src, idx)` on it. * If keyword `init` is provided, it is used to initialize the content of `dst`. Otherwise, the init values is inferred from the reduction operator `op` for some common operators (e.g. `init = 0` for `op = +`). * If `dstsize` is provided, it will be used to define the size of destination array, otherwise it will be inferred by `src` and `idx`. See [`scatter!`](/api/NN_Primitives/NNlib#NNlib.scatter!) for full details on how `idx` works. **Examples** ```julia julia> NNlib.scatter(+, [10,100,1000], [3,1,2]) 3-element Vector{Int64}: 100 1000 10 julia> NNlib.scatter(+, [1 2 3 4; 5 6 7 8], [2,1,1,5]) 2×5 Matrix{Int64}: 5 1 0 0 4 13 5 0 0 8 julia> NNlib.scatter(*, [10,200,3000], [1,4,2]; init = 10, dstsize = 6) 6-element Vector{Int64}: 100 30000 10 2000 10 10 ``` source ```julia NNlib.scatter!(op, dst, src, idx) ``` Scatter operation, which writes data in `src` into `dst` at locations `idx`. A binary reduction operator `op` is applied during the scatter. For each index `k` in `idx`, accumulates values in `dst` according to ```julia dst[:, ..., idx[k]...] = (op).(dst[:, ..., idx[k]...], src[:, ..., k...]) ``` See also [`scatter`](/api/NN_Primitives/NNlib#NNlib.scatter), [`gather`](/api/NN_Primitives/NNlib#NNlib.gather). **Arguments** * `op`: Operations to be applied on `dst` and `src`, e.g. `+`, `-`, `*`, `/`, `max`, `min` and `mean`. * `dst`: The destination for `src` to aggregate to. This argument will be mutated. * `src`: The source data for aggregating. * `idx`: The mapping for aggregation from source (index) to destination (value). The `idx` array can contain either integers or tuples. **Examples** ```julia julia> NNlib.scatter!(+, ones(3), [10,100], [1,3]) 3-element Vector{Float64}: 11.0 1.0 101.0 julia> NNlib.scatter!(*, fill(0.5, 2, 4), [1 10; 100 1000], [3,2]) 2×4 Matrix{Float64}: 0.5 5.0 0.5 0.5 0.5 500.0 50.0 0.5 ``` source ## Sampling {#Sampling} ````julia grid_sample(input::AbstractArray{T, 4}, grid::AbstractArray{T, 4}; padding_mode = :zeros) grid_sample(input::AbstractArray{T, 5}, grid::AbstractArray{T, 4}; padding_mode = :zeros) Given `input`, compute output by sampling `input` values at pixel locations from `grid`. Uses bilinear interpolation to calculate output values. This implementation assumes the extrema (`-1` and `1`) are considered as referring to the center points of the input’s corner pixels (i.e. align corners is `true`). # Arguments - `input`: Input array in `(W_in, H_in, [D_in,] C, N)` shape. - `grid`: Input grid in `(2, W_out, H_out, [D_out,] N)` shape. Where for each `(W_out, H_out, [D_out,] N)` grid contains `(x, y [,z])` coordinates that specify sampling locations normalized by the `input` shape. Therefore, `x`, `y` and [`z`] should have values in `[-1, 1]` range. For example, `(x = -1, y = -1, [z = -1])` is the left-top[-front] pixel of `input`, and `(x = 1, y = 1, [z = 1])` is the right-bottom-back pixel of `input`. Out-of-bound values are handled according to the `padding_mode`. - `padding_mode`: Out-of-bound padding. `:zeros` to use `0` for out-of-bound grid locations. `:border` to use border values for out-of-bound grid locations. Default is `:zeros`. # Returns `(W_out, H_out, [D_out,] C, N)` sampled grid from `input`. # Examples In the example below, grid contains two out-of-bound sampling locations, which are handled differently, depending on the `padding_mode`. ```jldoctest julia> x = reshape(collect(1.0:4.0), (2, 2, 1, 1)) 2×2×1×1 Array{Float64, 4}: [:, :, 1, 1] = 1.0 3.0 2.0 4.0 julia> grid = Array{Float64}(undef, 2, 3, 2, 1); julia> grid[:, 1, 1, 1] .= (-3, -1); julia> grid[:, 2, 1, 1] .= (0, -1); julia> grid[:, 3, 1, 1] .= (1, -1); julia> grid[:, 1, 2, 1] .= (-1, 1); julia> grid[:, 2, 2, 1] .= (0, 1); julia> grid[:, 3, 2, 1] .= (3, 1); julia> grid_sample(x, grid; padding_mode=:zeros) 3×2×1×1 Array{Float64, 4}: [:, :, 1, 1] = 0.0 3.0 1.5 3.5 2.0 0.0 julia> grid_sample(x, grid; padding_mode=:border) 3×2×1×1 Array{Float64, 4}: [:, :, 1, 1] = 1.0 3.0 1.5 3.5 2.0 4.0 ``` ```` source ```julia ∇grid_sample(Δ::AbstractArray{T, 4}, input::AbstractArray{T, 4}, grid::AbstractArray{T, 4}; padding_mode = :zeros) where T ``` **Arguments** * `Δ`: Input gradient in `(W_out, H_out, C, N)` shape (same as output of the primal computation). * `input`: Input from primal computation in `(W_in, H_in, C, N)` shape. * `grid`: Grid from primal computation in `(2, W_out, H_out, N)` shape. * `padding_mode`: Out-of-bound padding. `:zeros` to use `0` for out-of-bound grid locations. `:border` to use border values for out-of-bound grid locations. Should be the same as in primal computation. Default is `:zeros`. **Returns** `dinput` (same shape as `input`) and `dgrid` (same shape as `grid`) gradients. source ## Losses {#Losses} ```julia ctc_loss(ŷ, y) ``` Computes the connectionist temporal classification loss between `ŷ` and `y`. `ŷ` must be a classes-by-time matrices, i.e., each row represents a class and each column represents a time step. Additionally, the `logsoftmax` function will be applied to `ŷ`, so `ŷ` must be the raw activation values from the neural network and not, for example, the activations after being passed through a `softmax` activation function. `y` must be a 1D array of the labels associated with `ŷ`. The blank label is assumed to be the last label category in `ŷ`, so it is equivalent to `size(ŷ, 1)`. Used for sequence-to-sequence classification problems such as speech recognition and handwriting recognition where the exact time-alignment of the output (e.g., letters) is not needed to solve the problem. See [Graves et al. (2006)](https://www.cs.toronto.edu/~graves/icml_2006.pdf) or [Graves (2012)](https://www.cs.toronto.edu/~graves/preprint.pdf#chapter.7) for mathematical details. source ## Miscellaneous {#Miscellaneous} ```julia logsumexp(x; dims = :) ``` Computes `log.(sum(exp.(x); dims))` in a numerically stable way. Without `dims` keyword this returns a scalar. See also [`logsoftmax`](/api/NN_Primitives/NNlib#NNlib.logsoftmax). source ```julia glu(x, dim = 1) ``` The gated linear unit from the ["Language Modeling with Gated Convolutional Networks"](https://arxiv.org/abs/1612.08083) paper. Calculates `a .* sigmoid(b)`, where `x` is split in half along given dimension `dim` to form `a` and `b`. source ```julia @disallow_spawns ex ``` Disallow NNlib to use `@spawn` on divisible workloads. i.e. within `conv` etc. source ::: tip Tip `within_gradient` function currently doesn't work for Enzyme. Prefer to use `LuxLib.Utils.within_autodiff` if needed. Though pay heed that this function is not part of the public API. ::: ```julia within_gradient(x) --> Bool ``` Returns `false` except when used inside a `gradient` call, when it returns `true`. Useful for Flux regularisation layers which behave differently during training and inference. This should work with any ChainRules-based differentiation package, in which case `x` is ignored. But Tracker.jl overloads `with_gradient(x::TrackedArray)`, thus for widest use you should pass it an array whose gradient is of interest. There is also an overload for ForwardDiff.jl's `Dual` types (and arrays of them). **Examples** ```julia julia> using ForwardDiff, Zygote, NNlib julia> f_good(x) = if NNlib.within_gradient(x) @show 10x else x end; julia> Zygote.withgradient(f_good, 1.0) 10x = 10.0 (val = 10.0, grad = (10.0,)) julia> ForwardDiff.derivative(f_good, 1.0) 10x = Dual{ForwardDiff.Tag{typeof(f_good), Float64}}(10.0,10.0) 10.0 julia> f_bad(x, y) = if any(NNlib.within_gradient, (x, y)) @show x * y else x / y end; julia> Zygote.withgradient(f_bad, 2.0, 3.0) (val = 0.6666666666666666, grad = (0.3333333333333333, -0.2222222222222222)) julia> ForwardDiff.derivative(x -> f_bad(x, 3.0), 2.0) x * y = Dual{ForwardDiff.Tag{var"#9#10", Float64}}(6.0,3.0) 3.0 ``` What goes wrong in `f_bad` is that Zygote knows `any` to be non-differentiable, and thus completely ignores its contents. This is not a perfect mechanism, and the only style recommended is precisely that of `f_good` above. source ::: tip Tip Use `LuxLib.API.bias_activation!!` or `LuxLib.API.bias_activation` instead of `NNlib.bias_act!`. ::: ```julia bias_act!(σ, x, b) ``` This is equivalent to `x .= σ.(x .+ b)`, also replacing `sigmoid` & `tanh` with `sigmoid_fast` & `tanh_fast`. It will only overwrite `x` when `x isa StridedArray{<:AbstractFloat}`. When used within a gradient, it will overwrite only when `σ` has a method of `derivatives_given_output` which does not need the input at all. Such methods are defined by e.g. `@scalar_rule relu(x) Ω > 0` where the derivative contains only `Ω` (the output) not `x`. ::: warning Warning This is not safe to use if `x` is still needed for the gradient of some other function. Incorrect use will give silently wrong answers. It is intended mainly for Flux layers, in which the previous operation is known to be safe, e.g. `bias_act!(σ, weight * input, bias)` for a `Dense` layer. ::: source ## Dropout {#Dropout} ::: tip Tip Use `LuxLib.API.dropout` instead of `NNlib.dropout`. ::: ```julia dropout([rng], A, p; [dims]) ``` Returns an array in which each element of `A` is either replaced with zero, with probability `p`, or else multiplied by `1/(1-p)`. By default every element is treated independently. With keyword `dims=1`, a choice is made for every value of the 1st index i.e. each row of a matrix is either zero or not. Optional first argument is the random number generator used. **Examples** ```julia julia> dropout(ones(2, 10), 0.2) 2×10 Matrix{Float64}: 1.25 1.25 0.0 1.25 1.25 1.25 1.25 1.25 1.25 1.25 1.25 1.25 1.25 0.0 1.25 1.25 0.0 1.25 1.25 1.25 julia> mean(dropout(ones(10^4, 5), 0.2), dims=1) 1×5 Matrix{Float64}: 0.998 1.00075 0.99125 0.99575 1.00075 julia> dropout(ones(5, 5), 0.7, dims=1) # whole row the same 5×5 Matrix{Float64}: 3.33333 3.33333 3.33333 3.33333 3.33333 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 3.33333 3.33333 3.33333 3.33333 3.33333 0.0 0.0 0.0 0.0 0.0 julia> mean(dropout(ones(10^4, 5), 0.3, dims=1), dims=1) 1×5 Matrix{Float64}: 1.00571 1.00571 1.00571 1.00571 1.00571 ``` source ```julia dropout!(B, A, p; [dims]) ``` This does exactly `B .= dropout(A, p; dims)`, or rather, it's the implementation of out-of-place [`dropout`](/api/NN_Primitives/NNlib#NNlib.dropout). source ## Internal NNlib Functions {#Internal-NNlib-Functions} These functions are not part of the public API and are subject to change without notice. ```julia batched_transpose(A::AbstractArray{T,3}) batched_adjoint(A) ``` Equivalent to applying `transpose` or `adjoint` to each matrix `A[:,:,k]`. These exist to control how `batched_mul` behaves, as it operates on such matrix slices of an array with `ndims(A)==3`. `PermutedDimsArray(A, (2,1,3))` is equivalent to `batched_transpose(A)`, and is also understood by `batched_mul` (and more widely supported elsewhere). ```julia BatchedTranspose{T, S} <: AbstractBatchedMatrix{T, 3} BatchedAdjoint{T, S} ``` Lazy wrappers analogous to `Transpose` and `Adjoint`, returned by `batched_transpose` etc. source ```julia ∇conv_filter_direct!(dw, x, dy, cdims; alpha=1, beta=0) ``` Calculate the gradient imposed upon `w` in the convolution `y = x * w`. source ```julia _check_trivial_rotations!(out, arr, θ, rotation_center) ``` When `θ = 0 || π /2 || π || 3/2 || π` and if `rotation_center` is in the middle of the array. For an even array of size 4, the rotation\_center would need to be 2.5. For an odd array of size 5, the rotation\_center would need to be 3. In those cases, rotations are trivial just by reversing or swapping some axes. source ```julia NNlib.fast_act(f, [x::AbstractArray]) ``` Replaces `f == tanh` with [`tanh_fast`](/api/NN_Primitives/ActivationFunctions#NNlib.tanh_fast), etc. Takes an optional 2nd argument, so that you can disable this replacement for some array or element types. source ```julia spectrogram(waveform; pad::Int = 0, n_fft::Int, hop_length::Int, window, center::Bool = true, power::Real = 2.0, normalized::Bool = false, window_normalized::Bool = false, ) ``` Create a spectrogram or a batch of spectrograms from a raw audio signal. **Arguments** * `pad::Int`: Then amount of padding to apply on both sides. * `window_normalized::Bool`: Whether to normalize the waveform by the window’s L2 energy. * `power::Real`: Exponent for the magnitude spectrogram (must be ≥ 0) e.g., `1` for magnitude, `2` for power, etc. If `0`, complex spectrum is returned instead. See [`stft`](/api/NN_Primitives/NNlib#NNlib.stft) for other arguments. **Returns** Spectrogram in the shape `(T, F, B)`, where `T` is the number of window hops and `F = n_fft ÷ 2 + 1`. source ```julia is_strided(A::AbstractArray) -> Bool ``` This generalises `A isa StridedArray` to treat wrappers like `A::PermutedDimsArray`, for which it returns `is_strided(parent(A))`. It returns `true` for `CuArray`s, and `PermutedDimsArray`s of those. Other wrappers (defined outside Base, LinearAlgebra) are assumed not to break strided-ness, and hence also return `is_strided(parent(A))`. This correctly handles things like `NamedDimsArray` wihch don't alter indexing. However, it's a little pessimistic in that e.g. a `view` of such a container will return `false`, even in cases where the same `view` of `parent(A)` would be a `StridedArray`. source ```julia conv_direct!(y, x, w, cdims; alpha=1, beta=0) ``` Direct convolution implementation; used for debugging, tests, and mixing/matching of strange datatypes within a single convolution. Uses naive nested for loop implementation and does not attempt to optimize performance. Rather, this implementation is intended to be maximally understandable and debuggable, to aid in testing other, more performant implementations. We also explicitly support mixing and matching of strange datatypes, so that if the user really wants to convolve an image of `UInt8`'s with a `Float16` kernel, storing the result in a `Float32` output, there is at least a function call for that madness. The keyword arguments `alpha` and `beta` control accumulation behavior; this function calculates `y = alpha * x * w + beta * y`, therefore by setting `beta` to a nonzero value, the user is able to accumulate values into a preallocated `y` buffer, or by setting `alpha` to a nonunitary value, an arbitrary gain factor can be applied. By defaulting `beta` to `false`, we make use of the Bradbury promotion trick to override `NaN`'s that may pre-exist within our output buffer, as `false*NaN == 0.0`, whereas `0.0*NaN == NaN`. Only set `beta` if you are certain that none of the elements within `y` are `NaN`. The basic implementation performs 3-dimensional convolution; 1-dimensional and 2- dimensional cases are supported by simply reshaping `y`, `x` and `w`, for which wrapper methods are available. source ```julia gemm!() ``` Low-level gemm!() call with pointers, borrowed from Knet.jl Calculates `C = alpha*op(A)*op(B) + beta*C`, where: * `transA` and `transB` set `op(X)` to be either `identity()` or `transpose()` * alpha and beta are scalars * op(A) is an (M, K) matrix * op(B) is a (K, N) matrix * C is an (M, N) matrix. source ```julia calc_padding_regions(dims) ``` Padding is a jerk. A HUGE jerk that tries to sneak a bunch of conditionals and edge cases (quite literally) into our beautiful stencil operations such as convolution, pooling, etc... The way we deal with this is to, first, deal with everything in 3d, and then define a single padding region helper function that returns the seven regions that all 3d operations must deal with, including the central "unpadded" region where we can run at full bore, not paying any attention to padding. source ```julia ∇depthwiseconv_data_im2col!(dx, w, dy, cdims, col=similar(dx); alpha=1, beta=0) ``` Depwthwise conv2d backward pass onto the input using im2col and GEMM. See [`conv_im2col!`](/api/NN_Primitives/NNlib#NNlib.conv_im2col!) for explanation of optional parameters. source ```julia _prepare_imrotate(arr, θ, rotation_center) ``` Prepate `sin` and `cos`, creates the output array and converts type of `rotation_center` if required. source ```julia insert_singleton_spatial_dimension(cdims::ConvDims) ``` When converting a 1d convolution to a 2d, or a 2d to a 3d, we need to insert a singleton spatial dimension at the end of the spatial dimensions. This does so for a ConvDims. source ```julia _fast_broadcast!(f, x, y, z...) ``` This does `x .= f.(x, y, z...)`, but works around an issue with broadcasting that prevents SIMD in such cases. Can perhaps be removed once https://github.com/JuliaLang/julia/issues/43153 is fixed. Has an `rrule` to avoid mutation within derivatives. ::: warning Warning Not intended for general use. Uses `@inbounds` but does not check sizes! Assumes that `f` has no derivative! ::: source ```julia hann_window( window_length::Int, ::Type{T} = Float32; periodic::Bool = true, ) where T <: Real ``` Hann window function (ref: [Window function § Hann and Hamming windows - Wikipedia](https://en.wikipedia.org/wiki/Window_function#Hann_and_Hamming_windows)). $w\[n] = \frac{1}{2}\[1 - \cos(\frac{2 \pi n}{N - 1})]$ Where $N$ is the window length. ```julia julia> lineplot(hann_window(100); width=30, height=10) ┌──────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣠⠚⠉⠉⠉⠢⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡔⠁⠀⠀⠀⠀⠀⠘⢄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⠞⠀⠀⠀⠀⠀⠀⠀⠀⠈⢆⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⢀⡎⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⢣⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⡎⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⢦⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⢀⠞⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⢆⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⢀⡜⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⢇⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⢀⠎⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⢦⠀⠀⠀⠀│ │⠀⠀⠀⢠⠊⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠣⡀⠀⠀│ 0 │⣀⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⢤⣀│ └──────────────────────────────┘ ⠀0⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀100⠀ ``` **Arguments:** * `window_length::Int`: Size of the window. * `::Type{T}`: Elemet type of the window. **Keyword Arguments:** * `periodic::Bool`: If `true` (default), returns a window to be used as periodic function. If `false`, return a symmetric window. Following always holds: ```julia julia> N = 256; julia> hann_window(N; periodic=true) ≈ hann_window(N + 1; periodic=false)[1:end - 1] true julia> hann_window(N) ≈ hamming_window(N; α=0.5f0, β=0.5f0) true ``` **Returns:** Vector of length `window_length` and eltype `T`. source ```julia _rng_from_array(x) ``` Return the random number generator most appropriate for `x`: `CUDA.default_rng()` for `CuArray`, else `Random.default_rng()` source ```julia ∇depthwiseconv_filter_im2col!(dw, w, dy, cdims, col=similar(dw, ∇filter_im2col_dims(cdims)); alpha=1, beta=0) ``` Depthwise conv backward pass onto the weights using im2col and GEMM. See [`conv_im2col!`](/api/NN_Primitives/NNlib#NNlib.conv_im2col!) for explanation of optional parameters. source ```julia istft(y; n_fft::Int, hop_length::Int = n_fft ÷ 4, window = nothing, center::Bool = true, normalized::Bool = false, return_complex::Bool = false, original_length::Union{Nothing, Int} = nothing, ) ``` Inverse Short-time Fourier Transform. Return the least squares estimation of the original signal **Arguments:** * `y`: Input complex array in the `(n_fft, n_frames, B)` shape. Where `B` is the optional batch dimension. **Keyword Arguments:** * `n_fft::Int`: Size of Fourier transform. * `hop_length::Int`: Distance between neighboring sliding window frames. * `window`: Window function that was applied to the input of `stft`. If `nothing` (default), then no window was applied. * `center::Bool`: Whether input to `stft` was padded on both sides so that $t$-th frame is centered at time $t \times \text{hop length}$. Padding is done with `pad_reflect` function. * `normalized::Bool`: Whether input to `stft` was normalized. * `return_complex::Bool`: Whether the output should be complex, or if the input should be assumed to derive from a real signal and window. * `original_length::Union{Nothing, Int}`: Optional size of the first dimension of the input to `stft`. Helps restoring the exact `stft` input size. Otherwise, the array might be a bit shorter. source ```julia transpose_swapbatch(x::AbstractArray) ``` Given an AbstractArray, swap its batch and channel axes, as we must during transposed convolution. We do this to the operands during convolution, and then again to the output once we're done. source ```julia transpose_pad(cdims::ConvDims) ``` Transposed convolution can be calculated in terms of typical convolution with some extra padding. This method computes the padding of the convolution that would result in the transposed convolution of two operands, in essence taking care of that "extra padding". Note that this method should almost always be accompanied by a call that predilates one of the operands. source ```julia power_to_db(s; ref::Real = 1f0, amin::Real = 1f-10, top_db::Real = 80f0) ``` Convert a power spectrogram (amplitude squared) to decibel (dB) units. **Arguments** * `s`: Input power. * `ref`: Scalar w.r.t. which the input is scaled. * `amin`: Minimum threshold for `s`. * `top_db`: Threshold the output at `top_db` below the peak: `max.(s_db, maximum(s_db) - top_db)`. **Returns** `s_db ~= 10 * log10(s) - 10 * log10(ref)` source ```julia col2im!(x, col, cdims, beta=0) ``` Does the inverse of `im2col!()`, converting `col` back into a 3d image, used for backward passes, transposed convolutions, etc... Note that this method has not been optimized in the same way as `im2col()` has, because it is slightly more complicated due to the more chaotic data access patterns, and I'm not desperate enough yet. source ```julia depthwiseconv_im2col!(y, x, w, cdims, col=similar(x); alpha=1, beta=0) ``` Perform a depthwise convolution using im2col and GEMM, store the result in `y`. See [`conv_im2col!`](/api/NN_Primitives/NNlib#NNlib.conv_im2col!) for explanation of optional parameters. source ```julia storage_type(A) -> Type ``` Removes all wrappers to return the `Array` or `CuArray` (or whatever) type within. ```julia julia> view(reshape(ones(10)',2,5),:, 3:4) |> storage_type Array{Float64,1} julia> reshape(sparse(rand(10)), 5,2) |> storage_type SparseVector{Float64,Int64} ``` source ```julia im2col_dims(c::ConvDims) ``` im2col calculates, for each output pixel, the "convolution" of N kernels where N is the number of output channels, by doing a matrix multiply. The dimensions of that matrix are given by this function. Note that because im2col is multithreaded, we need to allocate a separate workspace of memory per-thread; hence the dimensions returned by this will depend on the number of threads Julia is currently running with. source ```julia ∇depthwiseconv_filter_direct!(dw, x, dy, cdims; alpha=1, beta=0) ``` Calculate the gradient imposed upon `w` in the depthwise convolution `y = x * w`. source ```julia reverse_indices(idx) ``` Return the reverse indices of `idx`. The indices of `idx` will be values, and values of `idx` will be index. **Arguments** * `idx`: The indices to be reversed. Accepts array or cuarray of integer, tuple or `CartesianIndex`. source ```julia ∇conv_filter_im2col!(dw, x, dy, cdims, col=similar(dw, ∇filter_im2col_dims(cdims)); alpha=1, beta=0) ``` Conv backward pass onto the weights using im2col and GEMM; stores the result in `dw`. See [`conv_im2col!`](/api/NN_Primitives/NNlib#NNlib.conv_im2col!) for explanation of optional parameters. source ```julia conv_im2col!(y, x, w, cdims, col=similar(x); alpha=1, beta=0) ``` Perform a convolution using im2col and GEMM, store the result in `y`. The kwargs `alpha` and `beta` control accumulation behavior; internally this operation is implemented as a matrix multiply that boils down to `y = alpha * x * w + beta * y`, thus by setting `beta` to a nonzero value, multiple results can be accumulated into `y`, or by setting `alpha` to a nonunitary value, various gain factors can be applied. Note for the particularly performance-minded, you can provide a pre-allocated `col`, which should eliminate any need for large allocations within this method. source ```julia ∇conv_data_direct!(dx, dy, w, cdims; alpha=1, beta=0) ``` Calculate the gradient imposed upon `x` in the convolution `y = x * w`. source Performs dimensional consistency checks and return the dimensionality of the scattered objects. source ```julia ∇conv_data_im2col!(dx, w, dy, cdims, col=similar(dx); alpha=1, beta=0) ``` Conv2d backward pass onto the input using im2col and GEMM; stores the result in `dx`. See [`conv_im2col!`](/api/NN_Primitives/NNlib#NNlib.conv_im2col!) for explanation of optional parameters. source ```julia storage_typejoin(A, B, C, ...) -> Type ``` Reduces with `Base.promote_typejoin`, in order that this conveys useful information for dispatching to BLAS. It does not tell you what container to allocate: ```julia julia> storage_typejoin(rand(2), rand(Float32, 2)) Array{T,1} where T julia> eltype(ans) <: LinearAlgebra.BlasFloat false julia> storage_typejoin(rand(2), rand(2,3), rand(2,3,4)) Array{Float64,N} where N ``` source ```julia add_blanks(z) ``` Adds blanks to the start and end of `z`, and between items in `z` source ```julia ∇filter_im2col_dims(c::ConvDims) ``` Like [`im2col_dims`](/api/NN_Primitives/NNlib#NNlib.im2col_dims), but saves some memory because multiple (Julia) threads are not required for the filter gradient calculation. Note: in the future, this may return `Dims{2}` instead of `Dims{3}`. source \_bilinear\_helper(yrot, xrot, yrot\_f, xrot\_f, yrot\_int, xrot\_int) Some helper variables source ```julia _triangular_filterbanks( freq_points::Vector{Float32}, all_freqs::Vector{Float32}) ``` Create triangular filter banks. **Arguments:** * `freq_points::Vector{Float32}`: Filter midpoints of size `n_filters`. * `all_freqs::Vector{Float32}`: Frequency points of size `n_freqs`. **Returns:** Array of size `(n_freqs, n_filters)`. source ```julia ∇depthwiseconv_data_direct!(dx, dy, w, cdims; alpha=1, beta=0) ``` Calculate the gradient imposed upon `x` in the depthwise convolution `y = x * w`. We make use of the fact that a depthwise convolution is equivalent to `C_in` separate normal convolutions between that channel of `x` and the `C_mult` different kernels that get applied to it. The output of such a convolution is the gradient imposed upon that particular channel of `x`, and so we simply walk through `x`, calculating the gradient for each batch and channel independently. source ```julia db_to_power(s_db; ref::Real = 1f0) ``` Inverse of [`power_to_db`](/api/NN_Primitives/NNlib#NNlib.power_to_db). source ```julia predilated_size(x_size::Tuple, dilation::Tuple) ``` Calculate the size of a predilated `x` given a particular dilation factor. This is used within `predilate()` and `transpose_cdims()`. source ```julia stft(x; n_fft::Int, hop_length::Int = n_fft ÷ 4, window = nothing, center::Bool = true, normalized::Bool = false, ) ``` Short-time Fourier transform (STFT). The STFT computes the Fourier transform of short overlapping windows of the input, giving frequency components of the signal as they change over time. $Y\[\omega, m] = \sum\_{k = 0}^{N - 1} \text{window}\[k] \text{input}\[m \times \text{hop length} + k] \exp(-j \frac{2 \pi \omega k}{\text{n fft}})$ where $N$ is the window length, $\omega$ is the frequency $0 \le \omega < \text{n fft}$ and $m$ is the index of the sliding window. **Arguments:** * `x`: Input, must be either a 1D time sequence (`(L,)` shape) or a 2D batch of time sequence (`(L, B)` shape). **Keyword Arguments:** * `n_fft::Int`: Size of Fourier transform. * `hop_length::Int`: Distance between neighboring sliding window frames. * `window`: Optional window function to apply. Must be 1D vector `0 < length(window) ≤ n_fft`. If window is shorter than `n_fft`, it is padded with zeros on both sides. If `nothing` (default), then no window is applied. * `center::Bool`: Whether to pad input on both sides so that $t$-th frame is centered at time $t \times \text{hop length}$. Padding is done with `pad_reflect` function. * `normalized::Bool`: Whether to return normalized STFT, i.e. multiplied with $\text{n fft}^{-0.5}$. **Returns:** Complex array of shape `(n_fft, n_frames, B)`, where `B` is the optional batch dimension. source ```julia hamming_window( window_length::Int, ::Type{T} = Float32; periodic::Bool = true, α::T = T(0.54), β::T = T(0.46), ) where T <: Real ``` Hamming window function (ref: [Window function § Hann and Hamming windows - Wikipedia](https://en.wikipedia.org/wiki/Window_function#Hann_and_Hamming_windows)). Generalized version of `hann_window`. $w\[n] = \alpha - \beta \cos(\frac{2 \pi n}{N - 1})$ Where $N$ is the window length. ```julia julia> lineplot(hamming_window(100); width=30, height=10) ┌──────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠚⠉⠉⠉⠢⡄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⠎⠁⠀⠀⠀⠀⠀⠈⢢⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⢠⠊⠀⠀⠀⠀⠀⠀⠀⠀⠀⢣⡀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⢰⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠱⡀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⣠⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠳⡀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⢰⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠱⡄⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⡰⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠱⡀⠀⠀⠀⠀│ │⠀⠀⠀⢀⠴⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠘⢄⠀⠀⠀│ │⠀⢀⡠⠊⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠳⣀⠀│ 0 │⠉⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉│ └──────────────────────────────┘ ⠀0⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀100⠀ ``` **Arguments:** * `window_length::Int`: Size of the window. * `::Type{T}`: Elemet type of the window. **Keyword Arguments:** * `periodic::Bool`: If `true` (default), returns a window to be used as periodic function. If `false`, return a symmetric window. Following always holds: ```julia julia> N = 256; julia> hamming_window(N; periodic=true) ≈ hamming_window(N + 1; periodic=false)[1:end - 1] true ``` * `α::Real`: Coefficient α in the equation above. * `β::Real`: Coefficient β in the equation above. **Returns:** Vector of length `window_length` and eltype `T`. source ```julia maximum_dims(dims) ``` Given an array of `CartesianIndex{N}` or `NTuple{N,Int}`, returns a tuple containing the maximum of all the 1st entries, all the 2nd entries, and so on up to `N`. Given an array of integers, returns `(maximum(dims),)`. (These arguments are what [`scatter`](/api/NN_Primitives/NNlib#NNlib.scatter) understands.) source ```julia batched_transpose(A::AbstractArray{T,3}) batched_adjoint(A) ``` Equivalent to applying `transpose` or `adjoint` to each matrix `A[:,:,k]`. These exist to control how `batched_mul` behaves, as it operates on such matrix slices of an array with `ndims(A)==3`. `PermutedDimsArray(A, (2,1,3))` is equivalent to `batched_transpose(A)`, and is also understood by `batched_mul` (and more widely supported elsewhere). ```julia BatchedTranspose{T, S} <: AbstractBatchedMatrix{T, 3} BatchedAdjoint{T, S} ``` Lazy wrappers analogous to `Transpose` and `Adjoint`, returned by `batched_transpose` etc. source ```julia _rotate_coordinates(sinθ, cosθ, i, j, rotation_center, round_or_floor) ``` This rotates the coordinates and either applies round(nearest neighbour) or floor for :bilinear interpolation) source ```julia melscale_filterbanks(; n_freqs::Int, n_mels::Int, sample_rate::Int, fmin::Float32 = 0f0, fmax::Float32 = Float32(sample_rate ÷ 2)) ``` Create triangular Mel scale filter banks (ref: [Mel scale - Wikipedia](https://en.wikipedia.org/wiki/Mel_scale)). Each column is a filterbank that highlights its own frequency. **Arguments:** * `n_freqs::Int`: Number of frequencies to highlight. * `n_mels::Int`: Number of mel filterbanks. * `sample_rate::Int`: Sample rate of the audio waveform. * `fmin::Float32`: Minimum frequency in Hz. * `fmax::Float32`: Maximum frequency in Hz. **Returns:** Filterbank matrix of shape `(n_freqs, n_mels)` where each column is a filterbank. ```julia julia> n_mels = 8; julia> fb = melscale_filterbanks(; n_freqs=200, n_mels, sample_rate=16000); julia> plot = lineplot(fb[:, 1]); julia> for i in 2:n_mels lineplot!(plot, fb[:, i]) end julia> plot ┌────────────────────────────────────────┐ 1 │⠀⡀⢸⠀⢸⠀⠀⣧⠀⠀⢸⡄⠀⠀⠀⣷⠀⠀⠀⠀⠀⣷⠀⠀⠀⠀⠀⠀⢀⣿⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⡇⢸⡆⢸⡇⠀⣿⠀⠀⡜⡇⠀⠀⢰⠋⡆⠀⠀⠀⢰⠁⡇⠀⠀⠀⠀⠀⡸⠀⢣⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⣿⢸⡇⡇⡇⢰⠹⡄⠀⡇⢱⠀⠀⢸⠀⢣⠀⠀⠀⡜⠀⢸⡀⠀⠀⠀⢀⠇⠀⠈⡇⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⣿⡇⡇⡇⡇⢸⠀⡇⢀⠇⠸⡀⠀⡇⠀⠸⡀⠀⢀⠇⠀⠀⢇⠀⠀⠀⡸⠀⠀⠀⠸⡄⠀⠀⠀⠀⠀⠀⠀│ │⢠⢻⡇⡇⡇⢱⢸⠀⢇⢸⠀⠀⡇⢀⠇⠀⠀⡇⠀⢸⠀⠀⠀⠸⡀⠀⢠⠇⠀⠀⠀⠀⢱⠀⠀⠀⠀⠀⠀⠀│ │⢸⢸⡇⢱⡇⢸⡇⠀⢸⢸⠀⠀⢣⢸⠀⠀⠀⢸⠀⡇⠀⠀⠀⠀⢇⠀⡜⠀⠀⠀⠀⠀⠈⢇⠀⠀⠀⠀⠀⠀│ │⢸⢸⡇⢸⠀⢸⡇⠀⢸⡇⠀⠀⢸⡎⠀⠀⠀⠈⣶⠁⠀⠀⠀⠀⠸⣤⠃⠀⠀⠀⠀⠀⠀⠘⡆⠀⠀⠀⠀⠀│ │⢸⠀⡇⢸⠀⠀⡇⠀⠀⡇⠀⠀⠀⡇⠀⠀⠀⠀⣿⠀⠀⠀⠀⠀⠀⣿⠀⠀⠀⠀⠀⠀⠀⠀⢱⡀⠀⠀⠀⠀│ │⢸⢸⡇⢸⠀⢸⡇⠀⢸⡇⠀⠀⢸⢇⠀⠀⠀⢀⠿⡀⠀⠀⠀⠀⢰⠛⡄⠀⠀⠀⠀⠀⠀⠀⠀⢣⠀⠀⠀⠀│ │⢸⢸⡇⡸⡇⢸⡇⠀⢸⢸⠀⠀⡜⢸⠀⠀⠀⢸⠀⡇⠀⠀⠀⠀⡎⠀⢣⠀⠀⠀⠀⠀⠀⠀⠀⠘⡆⠀⠀⠀│ │⢸⢸⡇⡇⡇⡸⢸⠀⡎⢸⠀⠀⡇⠈⡆⠀⠀⡇⠀⢸⠀⠀⠀⢰⠁⠀⠘⡆⠀⠀⠀⠀⠀⠀⠀⠀⠸⡄⠀⠀│ │⡇⢸⡇⡇⡇⡇⢸⠀⡇⠈⡆⢰⠁⠀⡇⠀⢰⠁⠀⠈⡆⠀⠀⡎⠀⠀⠀⢱⠀⠀⠀⠀⠀⠀⠀⠀⠀⢣⠀⠀│ │⡇⢸⢸⡇⡇⡇⠸⣰⠃⠀⡇⡸⠀⠀⢸⠀⡜⠀⠀⠀⢣⠀⢸⠁⠀⠀⠀⠈⡆⠀⠀⠀⠀⠀⠀⠀⠀⠈⢇⠀│ │⡇⡇⢸⠇⢸⡇⠀⣿⠀⠀⢣⡇⠀⠀⠸⣄⠇⠀⠀⠀⠸⡀⡇⠀⠀⠀⠀⠀⢱⠀⠀⠀⠀⠀⠀⠀⠀⠀⠸⡄│ 0 │⣇⣇⣸⣀⣸⣀⣀⣟⣀⣀⣸⣃⣀⣀⣀⣿⣀⣀⣀⣀⣀⣿⣀⣀⣀⣀⣀⣀⣈⣇⣀⣀⣀⣀⣀⣀⣀⣀⣀⣱│ └────────────────────────────────────────┘ ⠀0⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀200⠀ ``` source ```julia logaddexp(a, b) ``` Adds log-space `a` and `b` such that the result equals `log(exp(a)+exp(b))` source ```julia depthwiseconv_direct!(y, x, w, cdims; alpha=1, beta=0) ``` Direct depthwise convolution implementation; used for debugging, tests, and mixing/ matching of strange datatypes within a single convolution. Uses naive nested for loop implementation and does not attempt to optimize performance. Rather, this implementation is intended to be maximally understandable and debuggable, to aid in testing other, more performant implementations. We also explicitly support mixing and matching of strange datatypes, so that if the user really wants to convolve an image of `UInt8`'s with a `Float16` kernel, storing the result in a `Float32` output, there is at least a function call for that madness. One subtlety about depthwise convolutions; the shape of a depthwise convolutional kernel is `(spatial_dims..., C_mult, C_in)`, so the axis that must match with the number of channels in `x` is the last, not the second-to-last, as in a normal dense convolution. See the docstring for `conv_direct!()` for more on the optional parameters. source ```julia im2col!(col, x, cdims) ``` Converts a 3d image `x` into a matrix `col` for usage with GEMM-calculated convolution. Patches of `x` of size (kernel\_w, kernel\_h, kernel\_d, C\_in) will be extracted and laid out along the rows of `col`, one for each output pixel. This routine is used by all im2col-based convolutions, just with extra singleton dimensions added in the case of `2d` or `1d` images. source ```julia predilate(x, dilation::Tuple) ``` Places elements of `x` within a lattice of zeros, used in expressing a transposed convolution in terms of normal convolution. Note that while we call this "predilation" for aesthetic reasons, you are typically passing a "stride" value into here. Yes, transposed convolution is confusing. source ```julia safe_div(x, y) ``` Returns `x/y` unless `y==0`, in which case it just returns `x`. (Used internally by `scatter`.) source --- --- url: /dev/api/NN_Primitives/ActivationFunctions.md --- # Activation Functions {#NNlib-ActivationFunctions-API} Non-linearities that go between layers of your model. Note that, unless otherwise stated, activation functions operate on scalars. To apply them to an array you can call `σ.(xs)`, `relu.(xs)` and so on. ```julia celu(x, α=1) = x ≥ 0 ? x : α * (exp(x/α) - 1) ``` Activation function from ["Continuously Differentiable Exponential Linear Units"](https://arxiv.org/abs/1704.07483). ```julia julia> lineplot(celu, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉│ celu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡧⠶⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⠤⠔⠒⠋⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⠤⠤⠤⠤⠔⠒⠒⠒⠊⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> celu(-10f0) -0.9999546f0 ``` source ```julia elu(x, α=1) = x > 0 ? x : α * (exp(x) - 1) ``` Exponential Linear Unit activation function. See ["Fast and Accurate Deep Network Learning by Exponential Linear Units"](https://arxiv.org/abs/1511.07289). You can also specify the coefficient explicitly, e.g. `elu(x, 1)`. ```julia julia> lineplot(elu, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉│ elu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡧⠶⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⠤⠔⠒⠋⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⠤⠤⠤⠤⠔⠒⠒⠒⠊⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> elu(-10f0) -0.9999546f0 julia> elu(-10f0, 2) -1.9999092f0 ``` source ```julia gelu(x) = gelu_tanh(x) ``` Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415). See [`gelu_tanh`](/api/NN_Primitives/ActivationFunctions#NNlib.gelu_tanh). source ```julia hardσ(x) = max(0, min(1, (x + 3) / 6)) ``` Piecewise linear approximation of [`sigmoid`](/api/NN_Primitives/ActivationFunctions#NNlib.sigmoid). ```julia julia> lineplot(hardsigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋⠉⠉⠉⠉⠉⠉⠉⠉│ hardσ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡠⠔⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⡗⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠋⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⠤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot(sigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠒⠒⠋⠉⠉⠉⠉⠉⠉│ σ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⡏⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡔⠋⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⠤⠤⠤⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia hardσ(x) = max(0, min(1, (x + 3) / 6)) ``` Piecewise linear approximation of [`sigmoid`](/api/NN_Primitives/ActivationFunctions#NNlib.sigmoid). ```julia julia> lineplot(hardsigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋⠉⠉⠉⠉⠉⠉⠉⠉│ hardσ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡠⠔⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⡗⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠋⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⠤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot(sigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠒⠒⠋⠉⠉⠉⠉⠉⠉│ σ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⡏⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡔⠋⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⠤⠤⠤⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia sigmoid_fast(x) ``` This is a faster, and very slightly less accurate, version of `sigmoid`. For `x::Float32`, perhaps 3 times faster, and maximum errors 2 eps instead of 1. See also [`tanh_fast`](/api/NN_Primitives/ActivationFunctions#NNlib.tanh_fast). ```julia julia> sigmoid(0.2f0) 0.54983395f0 julia> sigmoid_fast(0.2f0) 0.54983395f0 julia> hardσ(0.2f0) 0.53333336f0 ``` source ```julia hardtanh(x) = max(-1, min(1, x)) ``` Segment-wise linear approximation of `tanh`, much cheaper to compute. See ["Large Scale Machine Learning"](https://ronan.collobert.com/pub/matos/2004_phdthesis_lip6.pdf). See also [`tanh_fast`](/api/NN_Primitives/ActivationFunctions#NNlib.tanh_fast). ```julia julia> lineplot(hardtanh, -2, 2, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⠔⠋⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│ hardtanh(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⢀⡤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡷⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠖⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠖⠋⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⠔⠋⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x julia> lineplot(tanh, -2, 2, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠤⠤⠒⠒⠒⠊⠉⠉⠉│ tanh(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⢀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⡷⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠔⠊⠁⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⣀⣀⣀⡠⠤⠤⠤⠖⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia tanh_fast(x) ``` This is a faster but slighly less accurate version of `tanh`. Where Julia's `tanh` function has an error under 2 eps, this may be wrong by 5 eps, a reduction by less than one decimal digit. For `x::Float32` this is usually about 10 times faster, with a smaller speedup for `x::Float64`. For any other number types, it just calls `tanh`. See also [`sigmoid_fast`](/api/NN_Primitives/ActivationFunctions#NNlib.sigmoid_fast). ```julia julia> tanh(0.5f0) 0.46211717f0 julia> tanh_fast(0.5f0) 0.46211714f0 julia> hard_tanh(0.5f0) 0.5f0 ``` source ```julia leakyrelu(x, a=0.01) = max(a*x, x) ``` Leaky [Rectified Linear Unit](https://en.wikipedia.org/wiki/Rectifier_\(neural_networks\)) activation function. You can also specify the coefficient explicitly, e.g. `leakyrelu(x, 0.01)`. ```julia julia> lineplot(x -> leakyrelu(x, 0.5), -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉│ #42(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠉⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⣤⣤⡤⡧⠶⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⠤⠤⠒⠒⠋⠉⠁⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⣀⣀⠤⠤⠒⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> leakyrelu(-10f0, 0.2) -2.0f0 julia> leakyrelu(-10f0, 0.02) -0.5f0 ``` source ```julia lisht(x) = x * tanh(x) ``` Activation function from ["LiSHT: Non-Parametric Linearly Scaled Hyperbolic Tangent ..."](https://arxiv.org/abs/1901.05894) ```julia julia> lineplot(lisht, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔│ lisht(x) │⠀⠈⠑⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀│ │⠀⠀⠀⠀⠈⠣⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠁⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠊⠁⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠢⡄⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⠔⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⢄⡀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⢀⡠⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⠦⣄⣀⣀⣇⣀⣀⠤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot!(ans, logcosh) ┌────────────────────────────────────────┐ 2 │⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔│ lisht(x) │⠀⠈⠑⢦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀│ logcosh(x) │⠢⣄⠀⠀⠈⠣⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠁⠀⠀⣀⠔│ f(x) │⠀⠈⠑⠢⣀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠊⠁⠀⣀⠔⠊⠁⠀│ │⠀⠀⠀⠀⠀⠉⠢⢄⡀⠉⠢⡄⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⠔⠋⠀⡠⠔⠋⠁⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⠦⣌⡓⢄⡀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⢀⡠⠖⣁⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⠪⠷⣦⣄⣀⣀⣇⣀⣀⣤⠶⠕⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia logcosh(x) ``` Return `log(cosh(x))` which is computed in a numerically stable way. ```julia julia> lineplot(logcosh, -5, 5, height=7) ┌────────────────────────────────────────┐ 5 │⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ logcosh(x) │⠉⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠋│ │⠀⠀⠀⠑⠢⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠑⠦⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠊⠁⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⠦⡀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⠦⡀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⢀⡤⠒⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠑⠢⢄⣀⣀⣇⣀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia logσ(x) ``` Return `log(σ(x))` which is computed in a numerically stable way. ```julia julia> lineplot(logsigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡧⠤⠔⠒⠒⠒⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│ logσ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠉⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⢀⡤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⣀⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⡤⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -6 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia logσ(x) ``` Return `log(σ(x))` which is computed in a numerically stable way. ```julia julia> lineplot(logsigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 0 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡧⠤⠔⠒⠒⠒⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│ logσ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠊⠉⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⢀⡤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⣀⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⡤⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -6 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia mish(x) = x * tanh(softplus(x)) ``` Activation function from ["Mish: A Self Regularized Non-Monotonic Neural Activation Function"](https://arxiv.org/abs/1908.08681). ```julia julia> lineplot(mish, -5, 5, height=7) ┌────────────────────────────────────────┐ 5 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋│ mish(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠒⠁⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⠔⠋⠁⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⢀⡠⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡤⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣧⣔⣊⣁⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀│ -1 │⠀⠀⠀⠀⠀⠀⠀⠀⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia relu(x) = max(0, x) ``` [Rectified Linear Unit](https://en.wikipedia.org/wiki/Rectifier_\(neural_networks\)) activation function. ```julia julia> lineplot(relu, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠋│ relu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠊⠁⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠊⠁⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⡠⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⡠⠖⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣇⠔⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia relu6(x) = min(max(0, x), 6) ``` [Rectified Linear Unit](https://en.wikipedia.org/wiki/Rectifier_\(neural_networks\)) activation function capped at 6. See ["Convolutional Deep Belief Networks"](https://www.cs.toronto.edu/~kriz/conv-cifar10-aug2010.pdf) from CIFAR-10. ```julia julia> lineplot(relu6, -10, 10, height=7) ┌────────────────────────────────────────┐ 6 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠎⠉⠉⠉⠉⠉⠉⠉⠉│ relu6(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⡔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⡤⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⡠⠎⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⠖⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡔⠃⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⡧⠋⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-10⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀10⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia rrelu(x, lo=1/8, hi=1/3) = max(a*x, x) # where `a` is randomly sampled from uniform distribution `U(lo, hi)` ``` Randomized Leaky Rectified Linear Unit activation function. See ["Empirical Evaluation of Rectified Activations"](https://arxiv.org/abs/1505.00853) You can also specify the bound explicitly, e.g. `rrelu(x, 0.0, 1.0)`. ```julia julia> lineplot(rrelu, -20, 10, height=7) ┌────────────────────────────────────────┐ 10 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠋│ rrelu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⢀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡤⠤⣤⣤⢤⣤⣤⠤⠤⠤⢼⠮⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⣰⢀⣆⡄⣄⡄⡠⡰⠦⠷⡜⢢⠷⠳⠢⠊⠉⠉⠀⠀⠁⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠃⠉⠙⠘⠃⠈⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -10 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-20⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀10⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> extrema(rrelu.(fill(-10f0, 1000))) (-3.3316886f0, -1.2548422f0) ``` source ```julia selu(x) = λ * (x ≥ 0 ? x : α * (exp(x) - 1)) λ ≈ 1.05070... α ≈ 1.67326... ``` Scaled exponential linear units. See ["Self-Normalizing Neural Networks"](https://arxiv.org/abs/1706.02515). ```julia julia> lineplot(selu, -3, 2, height=7) ┌────────────────────────────────────────┐ 3 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ selu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⠒│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⢀⣀⠤⠖⠊⠉⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⡠⠤⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⣉⠭⠛⡏⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⣀⡤⠤⠒⠊⠉⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -2 │⠤⠤⠖⠒⠒⠒⠒⠒⠒⠒⠉⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> selu(-10f0) -1.7580194f0 ``` source ```julia σ(x) = 1 / (1 + exp(-x)) ``` Classic [sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function) activation function. Unicode `σ` can be entered as `\sigma` then tab, in many editors. The ascii name `sigmoid` is also exported. See also [`sigmoid_fast`](/api/NN_Primitives/ActivationFunctions#NNlib.sigmoid_fast). ```julia julia> using UnicodePlots julia> lineplot(sigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠒⠒⠋⠉⠉⠉⠉⠉⠉│ σ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⡏⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡔⠋⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⠤⠤⠤⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> sigmoid === σ true ``` source ```julia σ(x) = 1 / (1 + exp(-x)) ``` Classic [sigmoid](https://en.wikipedia.org/wiki/Sigmoid_function) activation function. Unicode `σ` can be entered as `\sigma` then tab, in many editors. The ascii name `sigmoid` is also exported. See also [`sigmoid_fast`](/api/NN_Primitives/ActivationFunctions#NNlib.sigmoid_fast). ```julia julia> using UnicodePlots julia> lineplot(sigmoid, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠒⠒⠋⠉⠉⠉⠉⠉⠉│ σ(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣀⠔⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡠⡏⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡔⠋⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠊⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⠤⠤⠤⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> sigmoid === σ true ``` source ```julia softplus(x) = log(exp(x) + 1) ``` See ["Deep Sparse Rectifier Neural Networks"](http://proceedings.mlr.press/v15/glorot11a/glorot11a.pdf), JMLR 2011. ```julia julia> lineplot(softplus, -3, 3, height=7) ┌────────────────────────────────────────┐ 4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ softplus(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠁⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠔⠊⠁⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡠⠤⠒⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⡧⠤⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⡠⠤⠤⠤⠤⠔⠒⠒⠚⠉⠉⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot!(ans, relu) ┌────────────────────────────────────────┐ 4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ softplus(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣠│ relu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣠⡴⠞⠋⠁│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣤⡴⠞⠋⠁⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⣀⡠⢤⡲⠝⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⡧⠤⠒⠊⣉⠥⠚⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⣠⣤⣤⣤⣤⣔⣒⣒⣚⣉⣉⣁⣀⣇⠴⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> softplus(16f0) 16.0f0 ``` source ```julia softshrink(x, λ=0.5) = (x ≥ λ ? x - λ : (-λ ≥ x ? x + λ : 0)) ``` See ["Softshrink Activation Function"](https://www.gabormelli.com/RKB/Softshrink_Activation_Function). ```julia julia> lineplot(softshrink, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀│ softshrink(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⠉⠁│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⠒⠋⠁⠀⠀⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⣤⡤⠤⠤⠤⠤⠤⠤⡧⠤⠤⠤⠤⠶⠮⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⢀⣀⠤⠖⠒⠉⠁⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⣀⠤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -2 │⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot!(ans, tanhshrink) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀│ softshrink(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⣀⡤⠔⠒⣉⡡│ tanhshrink(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⣀⡤⠤⣒⣋⠥⠤⠒⠊⠉⠁⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⣤⣤⣤⣤⡤⠤⠤⠤⠤⠤⠤⡷⠶⠶⠶⠶⠶⠾⠿⠯⠭⠭⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⢀⣀⡠⠤⠖⢒⣋⠭⠗⠒⠉⠁⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠊⣉⠤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -2 │⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀ julia> softshrink.((-10f0, 10f0)) (-9.5f0, 9.5f0) ``` source ```julia softsign(x) = x / (1 + |x|) ``` See ["Quadratic Polynomials Learn Better Image Features"](http://www.iro.umontreal.ca/~lisa/publications2/index.php/attachments/single/205) (2009). ```julia julia> lineplot(softsign, -5, 5, height=7) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⣀⣀⣀⠤⠤⠤⠤⠤│ softsign(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⣀⡤⠖⠒⠋⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⡔⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡯⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠁⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⣀⠤⠤⠒⠋⠁⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⠒⠒⠒⠒⠒⠊⠉⠉⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot!(ans, tanh) ┌────────────────────────────────────────┐ 1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⢀⡤⠖⠊⠉⠉⠉⣉⣉⣉⣉⣉⠭⠭⠭⠭⠭│ softsign(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⡔⣃⡤⠖⠒⠋⠉⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ tanh(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣧⡞⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⡯⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡴⠃⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⣀⣀⠤⠤⠒⢋⠕⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⣒⣒⣒⣒⣒⣊⣉⣉⣉⣉⣁⣀⣀⡠⠤⠒⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> softsign(1f0) 0.5f0 julia> softsign(100f0) 0.990099f0 ``` source ```julia swish(x) = x * σ(x) ``` Self-gated activation function. See ["Swish: a Self-Gated Activation Function"](https://arxiv.org/abs/1710.05941). ```julia julia> lineplot(swish, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤│ swish(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠋⠁⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⢀⣀⡤⠔⠊⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⣤⣤⡤⡧⠴⠶⠯⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠉⠑⠒⠒⠒⠒⠒⠒⠒⠒⠒⠒⠉⠉⠉⠉⠁⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia hardswish(x) = x * hardσ(x) ``` Hard-Swish activation function. See ["Searching for MobileNetV3"](https://arxiv.org/abs/1905.02244). ```julia julia> lineplot(hardswish, -2, 5, height = 7) ┌────────────────────────────────────────┐ 5 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠔⠒⠉│ hardswish(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠔⠒⠉⠁⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠖⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⢀⣀⠤⠖⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣇⣤⣤⣖⣚⣉⣁⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀│ -1 │⠉⠒⠒⠒⠒⠉⠉⠉⠉⠁⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀5⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot(hardswish, -4, 0, height = 7); julia> lineplot!(ans, swish) ┌────────────────────────────────────────┐ 0 │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⢣⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡜│ hardswish(x) │⠒⠒⠢⠤⢄⣀⡀⠀⠀⠀⠀⠱⡄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⠎⠀│ swish(x) │⠀⠀⠀⠀⠀⠀⠈⠉⠑⠒⠦⢄⣘⢄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡴⠃⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠑⡖⠦⢄⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⢔⠏⠁⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠣⣄⠀⠉⠑⠒⠦⠤⢄⣀⣀⣀⣀⡠⠤⠖⣊⠕⠁⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠓⠤⡀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡤⠖⠁⠀⠀⠀⠀⠀⠀⠀│ -0.4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠉⠒⠢⠤⠤⠔⠒⠋⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-4⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀0⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> hardswish.(-5:5)' 1×11 adjoint(::Vector{Float64}) with eltype Float64: -0.0 -0.0 -0.0 -0.333333 -0.333333 0.0 0.666667 1.66667 3.0 4.0 5.0 ``` source ```julia tanhshrink(x) = x - tanh(x) ``` See ["Tanhshrink Activation Function"](https://www.gabormelli.com/RKB/Tanhshrink_Activation_Function). ```julia julia> lineplot(tanhshrink, -3, 3, height=7) ┌────────────────────────────────────────┐ 3 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ tanhshrink(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡠⠤⠖⠊│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⢀⣀⡠⠤⠒⠊⠉⠁⠀⠀⠀⠀│ f(x) │⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⢤⣤⡤⠤⠤⠤⠤⠤⠤⡷⠶⠶⠶⠶⠶⠮⠭⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⣀⡠⠴⠒⠊⠉⠁⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⡠⠴⠒⠊⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -3 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-3⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀3⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> tanhshrink.((-10f0, 10f0)) (-9.0f0, 9.0f0) ``` source ```julia trelu(x, theta=1) = x > theta ? x : 0 ``` Threshold gated rectified linear activation function. See ["Zero-bias autoencoders and the benefits of co-adapting features"](https://arxiv.org/abs/1402.3337) ```julia julia> lineplot(trelu, -2, 4, height=7) ┌────────────────────────────────────────┐ 4 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⡤⠖⠋│ trelu(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠖⠋⠁⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊⠁⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠴⠊⠁⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⣠⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⡏⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ 0 │⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣀⣇⣀⣀⣀⣀⣀⣀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀4⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia gelu_tanh(x) = 0.5x * (1 + tanh(√(2/π) * (x + 0.044715x^3))) ``` Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415) using tanh approximation. This implementation uses `tanh` which allows for better pattern matching and fusion in optimizing compilers compared to the sigmoid-based implementation. For a potentially faster implementation that uses `sigmoid_fast`, see [`gelu_sigmoid`](/api/NN_Primitives/ActivationFunctions#NNlib.gelu_sigmoid). ```julia julia> lineplot(gelu_tanh, -2, 2, height=7) ┌────────────────────────────────────────┐ 2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢀⡠⠔⠊│ gelu_tanh(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣀⠔⠊⠁⠀⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⣀⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⣀⡠⠤⠒⠉⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ │⣤⣤⣤⣤⣤⣤⣤⣤⡤⠤⠤⠤⠤⠤⠤⠤⣤⣤⣤⡤⡧⠶⠶⠭⠥⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤⠤│ │⠀⠀⠀⠀⠀⠀⠀⠀⠈⠉⠉⠉⠉⠉⠉⠉⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ -1 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-2⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀2⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ julia> lineplot(gelu_tanh, -5, 0, height=7); julia> lineplot!(ans, swish) ┌────────────────────────────────────────┐ 0 │⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠉⠒⠒⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢸│ gelu_tanh(x) │⠑⠒⠢⠤⣄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇│ swish(x) │⠀⠀⠀⠀⠀⠈⠉⠒⠤⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢆⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⣸⠁│ f(x) │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠒⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠑⢄⠀⠀⠀⠀⠀⠀⠀⠀⢠⡇⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠓⢄⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠓⣄⠀⠀⠀⠀⠀⢠⡞⠀⠀│ │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠉⠦⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠓⢄⣀⣀⡤⢣⠃⠀⠀│ -0.2 │⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠈⠓⣄⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⢠⠇⠀⠀⠀│ └────────────────────────────────────────┘ ⠀-5⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀0⠀ ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀x⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ ``` source ```julia gelu_sigmoid(x) = x * σ(√(8/π) * (x + 0.044715x^3)) ``` Alternative implementation of the GELU activation function using `sigmoid` instead of `tanh`. This is mathematically equivalent to [`gelu_tanh`](/api/NN_Primitives/ActivationFunctions#NNlib.gelu_tanh) but may be faster in some cases. The sigmoid-based implementation may prevent pattern matching and fusion in some optimizing compilers. Use [`gelu_tanh`](/api/NN_Primitives/ActivationFunctions#NNlib.gelu_tanh) if you need better compiler optimization support. See ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415). source ```julia gelu_erf(x) = xΦ(x) = 0.5x * (1 + erf(x/√2)) ``` Activation function from ["Gaussian Error Linear Units"](https://arxiv.org/abs/1606.08415) without approximation. The SpecialFunctions.jl package needs to be loaded to use this function. source --- --- url: /dev/api/Building_Blocks/LuxCore.md --- # LuxCore {#LuxCore} `LuxCore.jl` defines the abstract layers for Lux. Allows users to be compatible with the entirely of `Lux.jl` without having such a heavy dependency. If you are depending on `Lux.jl` directly, you do not need to depend on `LuxCore.jl` (all the functionality is exported via `Lux.jl`). ## Abstract Types {#Abstract-Types} ```julia abstract type AbstractLuxLayer ``` Abstract Type for all Lux Layers Users implementing their custom layer, **must** implement * `initialparameters(rng::AbstractRNG, layer::CustomAbstractLuxLayer)` – This returns a `NamedTuple` containing the trainable parameters for the layer. * `initialstates(rng::AbstractRNG, layer::CustomAbstractLuxLayer)` – This returns a NamedTuple containing the current state for the layer. For most layers this is typically empty. Layers that would potentially contain this include `BatchNorm`, `LSTM`, `GRU`, etc. Optionally: * `parameterlength(layer::CustomAbstractLuxLayer)` – These can be automatically calculated, but it is recommended that the user defines these. * `statelength(layer::CustomAbstractLuxLayer)` – These can be automatically calculated, but it is recommended that the user defines these. See also [`AbstractLuxContainerLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxContainerLayer) source ```julia abstract type AbstractLuxWrapperLayer{layer} <: AbstractLuxLayer ``` See [`AbstractLuxContainerLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxContainerLayer) for detailed documentation. This abstract type is very similar to [`AbstractLuxContainerLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxContainerLayer) except that it allows for a single layer to be wrapped in a container. Additionally, on calling [`initialparameters`](/api/Building_Blocks/LuxCore#LuxCore.initialparameters) and [`initialstates`](/api/Building_Blocks/LuxCore#LuxCore.initialstates), the parameters and states are **not** wrapped in a `NamedTuple` with the same name as the field. As a convenience, we define the fallback call `(::AbstractLuxWrapperLayer)(x, ps, st)`, which calls `getfield(x, layer)(x, ps, st)`. source ```julia abstract type AbstractLuxContainerLayer{layers} <: AbstractLuxLayer ``` Abstract Container Type for certain Lux Layers. `layers` is a tuple containing fieldnames for the layer, and constructs the parameters and states using those. Users implementing their custom layer can extend the same functions as in [`AbstractLuxLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxLayer). ::: tip Advanced Structure Manipulation Advanced structure manipulation of these layers post construction is possible via `Functors.fmap`. For a more flexible interface, we recommend using `Lux.Experimental.@layer_map`. ::: ::: tip `fmap` Support `fmap` support needs to be explicitly enabled by loading `Functors.jl` and `Setfield.jl`. ::: ::: warning Changes from Pre-1.0 Behavior Previously if `layers` was a singleton tuple, [`initialparameters`](/api/Building_Blocks/LuxCore#LuxCore.initialparameters) and [`initialstates`](/api/Building_Blocks/LuxCore#LuxCore.initialstates) would return the parameters and states for the single field `layers`. From `v1.0.0` onwards, even for singleton tuples, the parameters/states are wrapped in a `NamedTuple` with the same name as the field. See [`AbstractLuxWrapperLayer`](/api/Building_Blocks/LuxCore#LuxCore.AbstractLuxWrapperLayer) to replicate the previous behavior of singleton tuples. ::: source ## General {#General} ```julia apply(model, x, ps, st) ``` In most cases this function simply calls `model(x, ps, st)`. However, it is still recommended to call `apply` instead of `model(x, ps, st)` directly. Some of the reasons for this include: 1. For certain types of inputs `x`, we might want to perform preprocessing before calling `model`. For eg, if `x` is an Array of `ReverseDiff.TrackedReal`s this can cause significant regressions in `model(x, ps, st)` (since it won't hit any of the BLAS dispatches). In those cases, we would automatically convert `x` to a `ReverseDiff.TrackedArray`. 2. Certain user defined inputs need to be applied to specific layers but we want the datatype of propagate through all the layers (even unsupported ones). In these cases, we can unpack the input in `apply` and pass it to the appropriate layer and then repack it before returning. See the Lux manual on Custom Input Types for a motivating example. ::: tip Tip `apply` is integrated with `DispatchDoctor.jl` that allows automatic verification of type stability. By default this is "disable"d. For more information, see the [documentation](https://github.com/MilesCranmer/DispatchDoctor.jl). ::: source ```julia stateless_apply(model, x, ps) ``` Calls `apply` and only returns the first argument. This function requires that `model` has an empty state of `NamedTuple()`. Behavior of other kinds of models are undefined and it is the responsibility of the user to ensure that the model has an empty state. source ```julia check_fmap_condition(cond, tmatch::Union{Type, Nothing}, x) -> Bool ``` `fmap`s into the structure `x` and see if `cond` is satisfied for any of the leaf elements. **Arguments** * `cond` - A function that takes a single argument and returns a `Bool`. * `tmatch` - A shortcut to check if `x` is of type `tmatch`. Can be disabled by passing `nothing`. * `x` - The structure to check. **Returns** A Boolean Value source ```julia contains_lux_layer(l) -> Bool ``` Check if the structure `l` is a Lux AbstractLuxLayer or a container of such a layer. source ```julia display_name(layer::AbstractLuxLayer) ``` Printed Name of the `layer`. If the `layer` has a field `name` that is used, else the type name is used. source ```julia replicate(rng::AbstractRNG) ``` Creates a copy of the `rng` state depending on its type. source ```julia setup(rng::AbstractRNG, layer) ``` Shorthand for getting the parameters and states of the layer `l`. Is equivalent to `(initialparameters(rng, l), initialstates(rng, l))`. ::: warning Warning This function is not pure, it mutates `rng`. ::: source ## Parameters {#Parameters} ```julia initialparameters(rng::AbstractRNG, layer) ``` Generate the initial parameters of the layer `l`. source ```julia parameterlength(layer) ``` Return the total number of parameters of the layer `l`. source ## States {#States} ```julia initialstates(rng::AbstractRNG, layer) ``` Generate the initial states of the layer `l`. source ```julia statelength(layer) ``` Return the total number of states of the layer `l`. source ```julia testmode(st::NamedTuple) ``` Make all occurrences of `training` in state `st` – `Val(false)`. source ```julia trainmode(st::NamedTuple) ``` Make all occurrences of `training` in state `st` – `Val(true)`. source ```julia update_state(st::NamedTuple, key::Symbol, value; exclude=Internal.isleaf) ``` Recursively update all occurrences of the `key` in the state `st` with the `value`. `exclude` is a function that is passed to `Functors.fmap_with_path`'s `exclude` keyword. ::: warning Needs Functors.jl This function requires `Functors.jl` to be loaded. ::: source ```julia preserves_state_type(l::AbstractLuxLayer) ``` Return `true` if the layer `l` preserves the state type of the layer. For example, if the input state type for your layer is `T1` and the output state type is `T2`, then this function should return `T1 == T2`. By default this function returns `true` for `AbstractLuxLayer`. For container layers, this function returns `true` if all the layers in the container preserve the state type. source ## Layer size {#Layer-size} ```julia outputsize(layer, x, rng) ``` Return the output size of the layer. The fallback implementation of this function assumes the inputs were batched, i.e., if any of the outputs are Arrays, with `ndims(A) > 1`, it will return `size(A)[1:(end - 1)]`. If this behavior is undesirable, provide a custom `outputsize(layer, x, rng)` implementation). ::: warning Fallback Implementation The fallback implementation of this function is defined once `Lux.jl` is loaded. ::: ::: warning Changes from Pre-1.0 Behavior Previously it was possible to override this function by defining `outputsize(layer)`. However, this can potentially introduce a bug that is hard to bypass. See [this PR](https://github.com/LuxDL/LuxCore.jl/pull/43) for more information. ::: source ## Stateful Layer {#Stateful-Layer} ```julia StatefulLuxLayer{FT}(model, ps, st) StatefulLuxLayer(model, ps, st) ``` ::: warning Warning This is not a LuxCore.AbstractLuxLayer ::: ::: tip Tip This layer can be used as a drop-in replacement for `Flux.jl` layers. ::: A convenience wrapper over Lux layers which stores the parameters and states internally. This is meant to be used in internal implementation of layers. When using the definition of `StatefulLuxLayer` without `FT` specified, make sure that all of the layers in the model define [`LuxCore.preserves_state_type`](/api/Building_Blocks/LuxCore#LuxCore.preserves_state_type). Else we implicitly assume that the state type is preserved. **Usecases** * Internal implementation of `@compact` heavily uses this layer. * In SciML codebases where propagating state might involving [`Box`ing](https://github.com/JuliaLang/julia/issues/15276). For a motivating example, see the Neural ODE tutorial. * Facilitates Nested AD support in Lux. For more details on this feature, see the [Nested AD Manual Page](/manual/nested_autodiff#nested_autodiff). **Static Parameters** * If `FT = true` then the type of the `state` is fixed, i.e., `typeof(last(model(x, ps, st))) == st`. * If `FT = false` then type of the state might change. Note that while this works in all cases, it will introduce type instability. **Arguments** * `model`: A Lux layer * `ps`: The parameters of the layer. This can be set to `nothing`, if the user provides the parameters on function call * `st`: The state of the layer **Inputs** * `x`: The input to the layer * `ps`: The parameters of the layer. Optional, defaults to `s.ps` **Outputs** * `y`: The output of the layer source --- --- url: /dev/api/Building_Blocks/WeightInitializers.md --- # WeightInitializers {#WeightInitializers-API} This package is a light dependency providing common weight initialization schemes for deep learning models. ## Supported RNG Types {#Supported-RNG-Types-WeightInit} | **RNG Type / Package** | **Returned Array Type** | **Unsupported Functions** | | ---------------------------------:| -----------------------:| ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:| | `Random.jl` | `Array` | | | `StableRNGs.jl` | `Array` | | | `CUDA.CURAND.default_rng()` | `CuArray` | | | `CUDA.default_rng()` | `CuArray` | | | `GPUArrays.default_rng(CuArray)` | `CuArray` | | | `AMDGPU.rocrand_rng()` | `ROCArray` | | | `AMDGPU.gpuarrays_rng()` | `ROCArray` | | | `GPUArrays.default_rng(ROCArray)` | `ROCArray` | | | `Metal.gpuarrays_rng()` | `MtlArray` | [`orthogonal`](/api/Building_Blocks/WeightInitializers#WeightInitializers.orthogonal) | | `GPUArrays.default_rng(MtlArray)` | `MtlArray` | [`orthogonal`](/api/Building_Blocks/WeightInitializers#WeightInitializers.orthogonal) | | `oneAPI.gpuarrays_rng()` | `oneArray` | [`orthogonal`](/api/Building_Blocks/WeightInitializers#WeightInitializers.orthogonal), [`truncated_normal`](/api/Building_Blocks/WeightInitializers#WeightInitializers.truncated_normal) | | `GPUArrays.default_rng(oneArray)` | `oneArray` | [`orthogonal`](/api/Building_Blocks/WeightInitializers#WeightInitializers.orthogonal), [`truncated_normal`](/api/Building_Blocks/WeightInitializers#WeightInitializers.truncated_normal) | ## API Reference {#API-Reference} ### Main Functions {#Main-Functions} ```julia glorot_normal([::AbstractRNG=Utils.default_rng()], [T=Float32], size...; gain = 1) -> AbstractArray{T, length(size)} ``` Return an `AbstractArray{T}` of the given `size` containing random numbers drawn from a normal distribution with standard deviation `gain * sqrt(2 / (fan_in + fan_out))`. This method is described in \[1] and also known as Xavier initialization. **References** \[1] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." *Proceedings of the thirteenth international conference on artificial intelligence and statistics*. 2010. source ```julia glorot_uniform([::AbstractRNG=Utils.default_rng()], [T=Float32], size...; gain = 1) -> AbstractArray{T, length(size)} ``` Return an `AbstractArray{T}` of the given `size` containing random numbers drawn from a uniform distribution on the interval $\[-x, x]$, where `x = gain * sqrt(6 / (fan_in + fan_out))`. This method is described in \[1] and also known as Xavier initialization. **References** \[1] Glorot, Xavier, and Yoshua Bengio. "Understanding the difficulty of training deep feedforward neural networks." *Proceedings of the thirteenth international conference on artificial intelligence and statistics*. 2010. source ```julia identity_init([::AbstractRNG=Utils.default_rng()], [T=Float32], size...; gain::Number=1, shift::Union{Integer, Tuple{Integer, Integer}}=0) -> AbstractArray{T} ``` Constructs an array that aims to provide an identity mapping when used as parameters in most layers of a neural network. The identity mapping is scaled by the `gain` parameter. **Behavior** * 1D: Returns a `Vector` of zeros (useful for biases in layers where `input_size == output_size`). * 2D: Returns an identity matrix (useful for fully connected layers with equal input and output sizes). * More than 2D: Returns a tensor where the central slice along the last two dimensions is an identity matrix, and the rest are zeros (useful for convolutional layers, simulating an identity convolution). **Caveats** * Not all layers will result in an identity mapping when using this initializer. Exceptions include recurrent and normalization layers. * Layers must have `input_size == output_size` for a perfect identity mapping. In cases where this condition is not met, the function pads extra dimensions with zeros. * For convolutional layers to achieve an identity mapping, kernel sizes must be odd, and appropriate padding must be applied to ensure the output feature maps are the same size as the input feature maps. **Arguments** * `rng::AbstractRNG`: An optional random number generator, included for consistency with other initializers but ignored since the output is deterministic. * `T::Type{<:Number}`: The numeric type of the array elements. * `size...`: The dimensions of the array to be initialized. * `gain::Number=1`: A scaling factor applied to the identity mapping. * `shift::Union{Integer, Tuple{Integer, Integer}}=0`: An integer or a tuple specifying the circular shift applied to the output array. **Returns** * `AbstractArray{T}`: An array initialized to represent an identity mapping, scaled by `gain` and optionally shifted by `shift`. **Examples** ```julia julia> identity_init(Xoshiro(123), Float32, 5, 5) 5×5 Matrix{Float32}: 1.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 1.0 0.0 0.0 0.0 0.0 0.0 1.0 julia> identity_init(Xoshiro(123), Float32, 3, 3, 1, 1; gain=1.5) 3×3×1×1 Array{Float32, 4}: [:, :, 1, 1] = 0.0 0.0 0.0 0.0 1.5 0.0 0.0 0.0 0.0 ``` source ```julia kaiming_normal([::AbstractRNG=Utils.default_rng()], [T=Float32], size...; gain = √T(2)) -> AbstractArray{T, length(size)} ``` Return an `AbstractArray{T}` of the given `size` containing random numbers taken from a normal distribution standard deviation `gain / sqrt(fan_in)` **References** \[1] He, Kaiming, et al. "Delving deep into rectifiers: Surpassing human-level performance on imagenet classification." *Proceedings of the IEEE international conference on computer vision*. 2015. source ```julia kaiming_uniform([::AbstractRNG=Utils.default_rng()], [T=Float32], size...; gain = √T(2)) -> AbstractArray{T, length(size)} ``` Return an `AbstractArray{T}` of the given `size` containing random numbers drawn from a uniform distribution on the interval `[-x, x]`, where `x = gain * sqrt(3/fan_in)`. **References** \[1] He, Kaiming, et al. "Delving deep into rectifiers: Surpassing human-level performance on imagenet classification." *Proceedings of the IEEE international conference on computer vision*. 2015. source ```julia sparse_init([::AbstractRNG=Utils.default_rng()], [T=Float32], dims::Integer...; sparsity::Number, std::Number=0.01) -> AbstractArray{T} ``` Creates a sparsely initialized weight matrix with a specified proportion of zeroed elements, using random numbers drawn from a normal distribution for the non-zero elements. This method was introduced in \[1]. ::: tip Note The sparsity parameter controls the proportion of the matrix that will be zeroed. For example, a sparsity of 0.3 means that approximately 30% of the elements will be set to zero. The non-zero elements are distributed according to a normal distribution, scaled by the std parameter. ::: **Arguments** * `rng::AbstractRNG`: The random number generator to use. * `T::Type{<:Number}`: The numeric type of the elements in the returned array. * `dims::Integer...`: The dimensions of the weight matrix to be generated. * `sparsity::Number`: The proportion of elements to be zeroed. Must be between 0 and 1. * `std::Number=0.01`: The standard deviation of the normal distribution before applying `gain`. **Returns** * `AbstractArray{T}`: A sparsely initialized weight matrix of dimensions `dims` and type `T`. **Examples** ```julia julia> y = sparse_init(Xoshiro(123), Float32, 5, 5; sparsity=0.3, std=0.01); julia> y isa Matrix{Float32} true julia> size(y) == (5, 5) true ``` **References** \[1] Martens, J, "Deep learning via Hessian-free optimization" Proceedings of the 27th International Conference on International Conference on Machine Learning. 2010. source ```julia truncated_normal([::AbstractRNG=Utils.default_rng()], [T=Float32], size...; mean = 0, std = 1, lo = -2, hi = 2) -> AbstractArray{T, length(size)} ``` Return an `AbstractArray{T}` of the given `size` where each element is drawn from a truncated normal distribution. The numbers are distributed like `filter(x -> lo ≤ x ≤ hi, mean .+ std .* randn(100))`. source ```julia orthogonal([::AbstractRNG=Utils.default_rng()], [T=Float32], dims::Integer...; gain = 1) -> AbstractArray{T, length(dims)} ``` Return an `AbstractArray{T}` of the given dimensions (`dims`) which is a (semi) orthogonal matrix, as described in \[1]. The function constructs an orthogonal or semi-orthogonal matrix depending on the specified dimensions. For two dimensions, it returns a matrix where `dims = (rows, cols)`. For more than two dimensions, it computes an orthogonal matrix of size `prod(dims[1:(end - 1)])` by `dims[end]` before reshaping it to the original dimensions. Cannot construct a vector, i.e., `length(dims) == 1` is forbidden. **Arguments** * `rng::AbstractRNG`: Random number generator. * `T::Type{<:Real}`: The type of the elements in the array. * `dims::Integer...`: The dimensions of the array. * `gain::Number`: Scaling factor for the elements of the orthogonal matrix. **References** \[1] Saxe, McClelland, Ganguli. "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks", ICLR 2014, https://arxiv.org/abs/1312.6120 source ### Other Convenience Functions {#Other-Convenience-Functions} ::: warning Beware Unlike the other functions these ones don't take a type argument. ::: ```julia zeros16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float16, length(size)} ``` Return an `AbstractArray{Float16}` of the given `size` containing an AbstractArray of zeros. source ```julia ones16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float16, length(size)} ``` Return an `AbstractArray{Float16}` of the given `size` containing an AbstractArray of ones. source ```julia rand16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float16, length(size)} ``` Return an `AbstractArray{Float16}` of the given `size` containing random numbers from a uniform distribution. source ```julia randn16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float16, length(size)} ``` Return an `AbstractArray{Float16}` of the given `size` containing random numbers from a standard normal distribution. source ```julia zeros32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float32, length(size)} ``` Return an `AbstractArray{Float32}` of the given `size` containing an AbstractArray of zeros. source ```julia ones32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float32, length(size)} ``` Return an `AbstractArray{Float32}` of the given `size` containing an AbstractArray of ones. source ```julia rand32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float32, length(size)} ``` Return an `AbstractArray{Float32}` of the given `size` containing random numbers from a uniform distribution. source ```julia randn32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float32, length(size)} ``` Return an `AbstractArray{Float32}` of the given `size` containing random numbers from a standard normal distribution. source ```julia zeros64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float64, length(size)} ``` Return an `AbstractArray{Float64}` of the given `size` containing an AbstractArray of zeros. source ```julia ones64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float64, length(size)} ``` Return an `AbstractArray{Float64}` of the given `size` containing an AbstractArray of ones. source ```julia rand64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float64, length(size)} ``` Return an `AbstractArray{Float64}` of the given `size` containing random numbers from a uniform distribution. source ```julia randn64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{Float64, length(size)} ``` Return an `AbstractArray{Float64}` of the given `size` containing random numbers from a standard normal distribution. source ```julia zerosC16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF16, length(size)} ``` Return an `AbstractArray{ComplexF16}` of the given `size` containing an AbstractArray of zeros. source ```julia onesC16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF16, length(size)} ``` Return an `AbstractArray{ComplexF16}` of the given `size` containing an AbstractArray of ones. source ```julia randC16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF16, length(size)} ``` Return an `AbstractArray{ComplexF16}` of the given `size` containing random numbers from a uniform distribution. source ```julia randnC16([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF16, length(size)} ``` Return an `AbstractArray{ComplexF16}` of the given `size` containing random numbers from a standard normal distribution. source ```julia zerosC32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF32, length(size)} ``` Return an `AbstractArray{ComplexF32}` of the given `size` containing an AbstractArray of zeros. source ```julia onesC32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF32, length(size)} ``` Return an `AbstractArray{ComplexF32}` of the given `size` containing an AbstractArray of ones. source ```julia randC32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF32, length(size)} ``` Return an `AbstractArray{ComplexF32}` of the given `size` containing random numbers from a uniform distribution. source ```julia randnC32([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF32, length(size)} ``` Return an `AbstractArray{ComplexF32}` of the given `size` containing random numbers from a standard normal distribution. source ```julia zerosC64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF64, length(size)} ``` Return an `AbstractArray{ComplexF64}` of the given `size` containing an AbstractArray of zeros. source ```julia onesC64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF64, length(size)} ``` Return an `AbstractArray{ComplexF64}` of the given `size` containing an AbstractArray of ones. source ```julia randC64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF64, length(size)} ``` Return an `AbstractArray{ComplexF64}` of the given `size` containing random numbers from a uniform distribution. source ```julia randnC64([::AbstractRNG=Utils.default_rng()], size...; kwargs...) -> AbstractArray{ComplexF64, length(size)} ``` Return an `AbstractArray{ComplexF64}` of the given `size` containing random numbers from a standard normal distribution. source --- --- url: /dev/api/Testing_Functionality/LuxTestUtils.md --- # LuxTestUtils {#LuxTestUtils} ::: warning Warning This is a testing package. Hence, we don't use features like weak dependencies to reduce load times. It is recommended that you exclusively use this package for testing and not add a dependency to it in your main package Project.toml. ::: Implements utilities for testing **gradient correctness** and **dynamic dispatch** of Lux.jl models. ## Testing using JET.jl {#Testing-using-JET.jl} ```julia @jet f(args...) call_broken=false opt_broken=false ``` Run JET tests on the function `f` with the arguments `args...`. If `JET.jl` fails to compile, then the macro will be a no-op. **Keyword Arguments** * `call_broken`: Marks the test\_call as broken. * `opt_broken`: Marks the test\_opt as broken. All additional arguments will be forwarded to `JET.@test_call` and `JET.@test_opt`. ::: tip Tip Instead of specifying `target_modules` with every call, you can set global target modules using [`jet_target_modules!`](/api/Testing_Functionality/LuxTestUtils#LuxTestUtils.jet_target_modules!). ```julia using LuxTestUtils jet_target_modules!(["Lux", "LuxLib"]) # Expects Lux and LuxLib to be present in the module calling `@jet` ``` ::: **Example** ```julia julia> @jet sum([1, 2, 3]) target_modules=(Base, Core) Test Passed julia> @jet sum(1, 1) target_modules=(Base, Core) opt_broken=true call_broken=true Test Broken Expression: #= REPL[21]:1 =# JET.@test_opt target_modules = (Base, Core) sum(1, 1) ``` source ```julia jet_target_modules!(list::Vector{String}; force::Bool=false) ``` This sets `target_modules` for all JET tests when using [`@jet`](/api/Testing_Functionality/LuxTestUtils#LuxTestUtils.@jet). source ## Gradient Correctness {#Gradient-Correctness} ```julia test_gradients(f, args...; skip_backends=[], broken_backends=[], kwargs...) ``` Test the gradients of `f` with respect to `args` using the specified backends. The ground truth gradients are computed using FiniteDiff.jl (unless specified otherwise) on CPU. | Backend | ADType | CPU | GPU | Notes | |:-------------- |:------------------- |:--- |:--- |:----------------- | | Zygote.jl | `AutoZygote()` | ✔ | ✔ | | | Mooncake.jl | `AutoMooncake()` | ✔ | ✖ | | | ForwardDiff.jl | `AutoForwardDiff()` | ✔ | ✖ | `len ≤ 32` | | Enzyme.jl | `AutoEnzyme()` | ✔ | ✖ | Only Reverse Mode | **Arguments** * `f`: The function to test the gradients of. * `args`: The arguments to test the gradients of. Only `AbstractArray`s are considered for gradient computation. Gradients wrt all other arguments are assumed to be `NoTangent()`. **Keyword Arguments** * `skip_backends`: A list of backends to skip. * `broken_backends`: A list of backends to treat as broken. * `soft_fail`: If `true`, then the test will be recorded as a `soft_fail` test. This overrides any `broken` kwargs. Alternatively, a list of backends can be passed to `soft_fail` to allow soft\_fail tests for only those backends. * `enzyme_set_runtime_activity`: If `true`, then activate runtime activity for Enzyme. * `enable_enzyme_reverse_mode`: If `true`, then enable reverse mode for Enzyme. * `kwargs`: Additional keyword arguments to pass to `check_approx`. * `ground_truth_backend`: The backend to use for computing the ground truth gradients. Defaults to `AutoFiniteDiff()`. * `ground_truth_eltype`: The eltype to use for computing the ground truth gradients. Defaults to `Float64`. **Example** ```julia julia> f(x, y, z) = x .+ sum(abs2, y.t) + sum(y.x.z) julia> x = (; t=rand(10), x=(z=[2.0],)) julia> test_gradients(f, 1.0, x, nothing) ``` source ```julia @test_gradients(f, args...; kwargs...) ``` See the documentation of [`test_gradients`](/api/Testing_Functionality/LuxTestUtils#LuxTestUtils.test_gradients) for more details. This macro provides correct line information for the failing tests. source ## Extensions to `@test` {#Extensions-to-@test} ```julia @test_softfail expr ``` Evaluate `expr` and record a test result. If `expr` throws an exception, the test result will be recorded as an error. If `expr` returns a value, and it is not a boolean, the test result will be recorded as an error. If the test result is false then the test will be recorded as a broken test, else it will be recorded as a pass. source ## Private API {#Private-API} ```julia mooncake_gradient_function(f, x) ``` Compute gradient using Mooncake.jl's value\_and\_gradient!! function. Returns only the gradient for args x. source --- --- url: /dev/tutorials/intermediate/2_BayesianNN.md --- ::: danger Not Run on CI This tutorial is not run on CI to reduce the computational burden. If you encounter any issues, please open an issue on the [Lux.jl](https://github.com/LuxDL/Lux.jl) repository. ::: # Bayesian Neural Network {#Bayesian-Neural-Network} This tutorial has been upstreamed from Lux to the official Turing Documentation. See https://turinglang.org/docs/tutorials/bayesian-neural-networks/ for the updated version. *** *This page was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).* --- --- url: /dev/references.md --- # References {#References} 1. A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser and I. Polosukhin. *Attention is all you need*. Advances in neural information processing systems **30** (2017). 2. J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo and Y. Liu. *Roformer: Enhanced transformer with rotary position embedding*. Neurocomputing **568**, 127063 (2024). 3. I. Goodfellow, D. Warde-Farley, M. Mirza, A. Courville and Y. Bengio. *Maxout networks*. In: *International conference on machine learning* (PMLR, 2013); pp. 1319–1327. 4. D. Ulyanov, A. Vedaldi and V. Lempitsky. *Instance normalization: The missing ingredient for fast stylization*, arXiv preprint arXiv:1607.08022 (2016). 5. T.-Y. Lin, P. Goyal, R. Girshick, K. He and P. Dollár. *Focal loss for dense object detection*. In: *Proceedings of the IEEE international conference on computer vision* (2017); pp. 2980–2988. 6. F. Milletari, N. Navab and S.-A. Ahmadi. *V-net: Fully convolutional neural networks for volumetric medical image segmentation*. In: *2016 fourth international conference on 3D vision (3DV)* (Ieee, 2016); pp. 565–571. 7. R. Hadsell, S. Chopra and Y. LeCun. *Dimensionality reduction by learning an invariant mapping*. In: *2006 IEEE computer society conference on computer vision and pattern recognition (CVPR'06)*, Vol. 2 (IEEE, 2006); pp. 1735–1742. 8. G. Klambauer, T. Unterthiner, A. Mayr and S. Hochreiter. *Self-normalizing neural networks*. Advances in neural information processing systems **30** (2017). 9. N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever and R. Salakhutdinov. *Dropout: a simple way to prevent neural networks from overfitting*. The journal of machine learning research **15**, 1929–1958 (2014). 10. S. Ioffe and C. Szegedy. *Batch normalization: Accelerating deep network training by reducing internal covariate shift*. In: *International conference on machine learning* (pmlr, 2015); pp. 448–456. 11. Y. Wu and K. He. *Group normalization*. In: *Proceedings of the European conference on computer vision (ECCV)* (2018); pp. 3–19. 12. J. L. Ba, J. R. Kiros and G. E. Hinton. *Layer normalization*, arXiv preprint arXiv:1607.06450 (2016).