Lux.jl Deep Learning Library

Lux is a neural network framework built using pure functions to ensure compatibility with compilers and automatic differentiation (AD) systems. The core mental model relies on the following principles:

• Immutability: Layers cannot store parameters or state internally. Instead, they store only the information required to construct them.
• Pure Functions: Layers are implemented as pure functions. Given the same inputs, parameters, and state, they must always produce the same output.
• State Management: Layers return a Tuple containing the result and the updated state. This includes handling stochasticity (e.g., Lux.Dropout) by passing and updating rngs within the state.
• Separation of Concerns: Lux explicitly separates layer structures from parameter data and state variables. This makes implementing techniques like WeightNorm or SpectralNorm trivial and avoids the issues found in frameworks that couple models with parameters.

Lux and Flux follow different design philosophies for custom layers:

1. Architecture vs. Data: In Lux, the layer struct defines the architecture, while the data (trainable parameters and non-trainable states) is stored separately in ps and st.
2. Avoid Mutables in Layers: Do not store mutable structures like Arrays directly inside a Lux layer struct. Instead, store initialization functions (e.g., () -> copy(A)) and use them within Lux.initialparameters and Lux.initialstates to perform lazy initialization.
3. Distinguishing Parameters and States:
   • Parameters (ps): Trainable values (e.g., weights, biases). Defined via Lux.initialparameters.
   • States (st): Non-trainable values (e.g., running means in BatchNorm, or fixed constants). Defined via Lux.initialstates.
4. Device Transfer: Since data is stored in ps and st, device transfer utilities (like gpu_device) must be applied to the parameters and states, not the layer itself.

Example Implementation: To implement a layer computing $A imes B imes x$ where $A$ is a non-trainable state and $B$ is a trainable parameter:

struct LuxLinear <: Lux.AbstractLuxLayer
    init_A
    init_B
end

# Use functions for lazy initialization to avoid storing arrays in the struct
function LuxLinear(A::AbstractArray, B::AbstractArray)
    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(),)

# Forward pass signature
(l::LuxLinear)(x, ps, st) = st.A * ps.B * x, st

using Lux, Random, NNlib, Zygote

struct LuxLinear <: Lux.AbstractLuxLayer
    init_A
    init_B
end

function LuxLinear(A::AbstractArray, B::AbstractArray)
    return LuxLinear(() -> copy(A), () -> copy(B))
end

Lux.initialparameters(::AbstractRNG, layer::LuxLinear) = (B=layer.init_B(),)
Lux.initialstates(::AbstractRNG, layer::LuxLinear) = (A=layer.init_A(),)

(l::LuxLinear)(x, ps, st) = st.A * ps.B * x, st