Neural network backends

HydroModels keeps neural-network integrations optional. The existing Lux backend remains the full-featured backend; Flux and SimpleChains are loaded only when their packages are loaded.

BackendForwardExplicit parametersExplicit stateMooncakeRecommended use
Luxyesyesyesproject-level validationstateful and general networks
FluxyesFlux.destructureno in this extensionfuture candidatestateless Dense/Chain models
SimpleChainsyesnative flat vectornonolightweight CPU forward evaluation

Flux

Load Flux before constructing a Flux-backed component:

using HydroModels
using Flux

chain = Flux.Chain(Flux.Dense(2 => 8, tanh), Flux.Dense(8 => 1))
flux = NeuralFlux(inputs, outputs, chain; chain_name=:runoff_net)

The extension calls Flux.destructure once when the component is created. The current model parameters are stored under nns.<chain_name>.params as a ComponentVector. Each nn_func(x, p) call rebuilds a runtime model from the provided parameter vector, so parameter changes are never hidden by a stale cache. This is correct for calibration but can allocate in a high-frequency ODE RHS.

The supported first-stage subset is stateless Flux.Dense and Flux.Chain compositions using ordinary pure forward activations such as identity, relu, tanh, and sigmoid. BatchNorm, Dropout, recurrent layers, GPU execution, and custom mutable layers are not part of this compatibility promise.

SimpleChains

using HydroModels
using SimpleChains

chain = SimpleChain(
    SimpleChains.static(2),
    TurboDense(tanh, 8),
    TurboDense(identity, 1),
)
flux = NeuralFlux(inputs, outputs, chain; chain_name=:runoff_net)

SimpleChains is natively structure/parameter separated. HydroModels calls SimpleChains.init_params(chain, Float32; rng=...) during parameter initialization and stores the resulting flat vector under params. The backend supports the feed-forward SimpleChain/TurboDense subset used by the HydroModels factories and is CPU forward-only.

SimpleChains backend is forward-only in HydroModels. Direct Mooncake differentiation is unsupported because the tested SimpleChains execution path uses llvmcall-based kernels that Mooncake cannot currently translate. This is a tested NO-GO for end-to-end Mooncake training, not an untested feature.

Factories select the backend explicitly:

create_neural_bucket(Val(:flux); ...)
create_simple_neural_bucket(Val(:flux); ...)
create_neural_bucket(Val(:simplechains); ...)
create_simple_neural_bucket(Val(:simplechains); ...)

The default no-argument factory continues to select Lux when Lux is loaded. No Flux or SimpleChains package is loaded by using HydroModels.