Speed up HybridModel without changing the user API (#3)

[BernhardAhrens]

• Make HybridModel forward pass type-stable to cut Zygote gradient overhead

The HybridModel forward pass returned a type-unstable output (out::Any), which forced Zygote to build slow, boxed pullbacks. Two causes:

1. The global/neural/fixed parameter NamedTuples were constructed from the model's Vector{Symbol} name fields at runtime, so their keys were unknown to the compiler (abstract NamedTuples).
2. _mechanistic_kwargs/_extra_params introspected the mechanistic model via methods(f)/kwarg_decl on every forward pass and returned a runtime tuple, making the mechanistic-model call infer as Any.

Encode the parameter-name groups (NP/GP/FP), the mechanistic kwarg plan (KW accepted kwargs, EX surfaced extras) and the mechanistic-function type (MM) as HybridModel type parameters, computed once in the constructor. Build the parameter NamedTuples from the type-level name tuples and select kwargs with NamedTuple{KW}(all_kwargs). The per-call methods(f) reflection now runs once at construction instead of on the hot path.

Result on an identical NN with an x->x^2 head (batch 512): forward overhead vs plain Lux drops from ~1.55x to ~1.09x and reverse-mode gradient overhead from ~4.3x to ~2.2x. Public API and behavior (including loss-only params) are unchanged; full test suite passes.

• Make training-loss assembly type-stable to speed up the train step

Profiling the per-minibatch training step showed the loss machinery was the next bottleneck after the forward pass: assemble_loss/_compute_loss returned Body::Any, so the differentiated loss value was type-unstable and Zygote built a slow, boxed pullback. Causes:

• assemble_loss iterated the runtime targets::Vector{Symbol} and indexed the prediction/target NamedTuples by a runtime Symbol, producing a Vector{Any}.
• The extra-loss branch (ExtraLoss.f::Union{Function,Nothing}, and the abstractly-typed LoggingLoss.extra_loss field) widened the loss value to Any even when no extra loss was configured.
• The training-loss symbol was passed as a runtime value, so loss_fn(..., Val(sym)) was a dynamic dispatch.

Changes (behavior-preserving; the generic vector-based _compute_loss and assemble_loss used by the unit tests / eval mode are unchanged):

• Carry target names in the HybridModel type domain (new TG type parameter) and add a type-stable _train_loss/_assemble_loss fast path that maps over the compile-time target tuple. Other model types fall back to the generic path.
• Parameterize ExtraLoss{F} and LoggingLoss{L,E,T} so the "no extra loss" case is a concretely-typed no-op (dispatched via _add_extra_train_loss/eval).
• Parameterize SymbolicLoss{S} and add _static_loss_spec so the training-loss symbol is a compile-time Val, giving a static loss_fn dispatch.

Full single_train_step! (RbQ10-style model, AutoZygote) improves ~12% at batch 512 (706->623 us) and ~14% at batch 256, on top of the forward-pass fix. All 28 testsets pass.

• Make parameter scaling type-stable via typed ParameterContainer and Val names

The remaining HybridModel Any (global_params / y_pred element types) came from scale_single_param slicing a ComponentArray by a runtime name, and from parameters::ParameterContainer hiding the concrete values NamedTuple.

Three approaches were prototyped (see exploration notes):

1. Type-parameterize HybridModel on the ParameterContainer (PC) and scale via getfield(pc.values, N) with Val{name}. Also lift scale_nn_outputs into the type domain so the scale/no-scale branch is not a Union. Fully concrete y_pred; -22% gradient vs previous.
2. Cache lowers/uppers as NamedTuples of scalars. Same concreteness and speed, but adds fields. Not taken.
3. @generated mechanistic apply that unrolls f(; a=nt.a, ...). Alone it does not make y_pred concrete (the splat is fine once kwargs are a concrete NamedTuple). Combined with (2) is only ~5 us faster. Not taken.

This commit is approach 1. Public scale_single_param(::Symbol, ...) and constructHybridModel are unchanged. 28/28 tests pass.

• Skip AD tracing of keyword-argument sorting to speed up gradients

EasyHybrid differentiates through keyword calls on its hot path: the mechanistic model is invoked as f(; forcings..., params...) and the training loss goes through compute_loss. Julia implements a keyword call by splitting the caller's NamedTuple into the callee's declared and defaulted arguments, using the Base helpers diff_names/sym_in/merge_names.

Those helpers return Bools and tuples of Symbols, so they hold no derivative information, but Zygote does not know that and builds an uninferred pullback for each of them on every keyword call. Profiling a gradient through a mechanistic model showed that single uninferred pullback taking ~24% of the samples, making a keyword call ~3x more expensive than the equivalent positional one (60.6µs vs 19.0µs at batch 512).

Declaring them non-differentiable removes the pullbacks. Gradient values are bit-identical before and after; only wasted work is dropped. Measured on an RbQ10-style model, the gradient gets 30%/23%/10% faster at batch 32/512/4096.

The helpers are Base internals, so each rule is guarded on the function still existing to stay safe across Julia versions.

• Pass logging positionally on the differentiated training paths

Both training drivers built their loss closure with a keyword call, compute_loss(model, ps, st, data; logging = logging), so Zygote had to trace Julia's keyword-splitting bookkeeping once per gradient evaluation.

compute_loss now takes logging positionally, with the keyword form kept as a thin wrapper so the public API and existing call sites are unchanged. The two training closures use the positional form.

Worth ~3% of the train step at batch 32 on top of the non-differentiable keyword-sorting rules; gradient values are unchanged.

• Read HM.targets only where it is used

compute_loss read HM.targets unconditionally, but only the ParamLoss branch and evaluation mode use it: on the training hot path the per-target machinery takes the target names from the model type (TG) instead.

Reading the field makes Zygote trace a getproperty on the model struct, which cost ~23µs per gradient (about 6% at batch 512) for a value that was then discarded. Moving the read into the two branches that need it removes that.

• Keep parameters as a NamedTuple unless the AD backend needs an array

init_model_state wrapped ps in a ComponentArray for every optimizer path. That conversion arrived alongside ForwardDiff support, which genuinely needs a flat vector, but it also landed on the default Zygote + Optimisers path, where it costs performance: inside a ComponentArray the dense layer weights are reshaped views into one flat vector rather than plain Matrixes, so LuxLib cannot take its fast matmul path.

Convert only for backends that require an AbstractArray. AutoForwardDiff is the only one in Lux that does; the reverse-mode backends walk the parameter tree with fmap and accept either container.

A ComponentArray restored via train_from is left as-is: unwrapping it with NamedTuple hands back the reshaped views rather than materializing them, so it would keep the slow path anyway.

• Document the parameter container change

The opt docstring stated that ps is wrapped in a ComponentArray on the single-rule Optimisers path, which is no longer true by default.

• Specialize the training loss path and drop mechanistic kwarg splat

Make train_mode a LoggingLoss type parameter so Zygote only sees the training return type, and call the mechanistic model with generated explicit keywords instead of merge + splat.

• Generate HybridModel parameter assembly instead of runtime merges

Unroll NN slicing, global/fixed scaling, and output packing from the existing type parameters so Zygote does not trace intermediate NamedTuples.

• Fuse the Vector HybridModel forward and skip packing in training

Collapse NN scaling, parameter binding, and the mechanistic call into one generated function so Zygote records fewer wrappers. Training MSE uses the mechanistic output directly, drops state from the AD tape, and avoids boolean gathering when the NaN mask is a Bool array.

• Fix generated HybridModel destructuring of the neural-net output

:(nn_out, st_nn = apply(...)) is a tuple expression in Julia, so nn_out was looked up as a global. Assign to (nn_out, st_nn) explicitly.

• Skip parameter packing on the training loss path

Training MSE only needs mechanistic predictions, so the generated Vector forward can omit the parameters NamedTuple. Revert boolean ifelse MSE; boolean gather is faster on the all-true masks used in training.

• Add synthetic RbQ10 performance regression test

Guard the fused Vector training loss against a 10% slowdown versus a Lux same-math gradient on the synthetic respiration example.

---

[BernhardAhrens]

check where merge_names is used

[BernhardAhrens]

put into one place

[BernhardAhrens]

what does generated do

[lazarusA]

done in #307