diff --git a/CHANGELOG.md b/CHANGELOG.md index dfc778ae..f6cfa3d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog ## Unreleased +- Faster default `HybridModel` training without changing the user API. Mechanistic models stay keyword functions; the public `HybridModel` call still returns predictions, `parameters`, extras, and `nn_outputs`. On the synthetic RbQ10 example from the docs (neural `:rb`, global `:Q10`, hidden `[16, 16]`, sigmoid, `scale_nn_outputs=true`, `input_batchnorm=true`, batch 512) the training-loss gradient is about 3× faster than before (~1007 µs → ~315 µs) and about 1.17× a hand-written Lux model of the same math (was ~3×). `train` on 20 000 rows, 20 epochs, batch 512 goes from ~1.55 s to ~0.45 s. Only the default Vector `:mse` path (no `extra_loss`) is fused; the public forward is unchanged. +- Default `train` keeps `NamedTuple` parameters from `Lux.setup`; a `ComponentArray` is used only for `AutoForwardDiff` and `Optimization.jl`. Whole-vector ops on `out.ps` (`length`, `norm`, `Array`) need `ComponentArray(ps)` first. ## v0.2.1 - 2026-09-04 - Set Runic version to 1 in tools in order to match the one in CI. This should take care of the regular updates for patches, locally as well as remote. Older Manifest files in formatter should be updated always accordingly [#299](https://github.com/EarthyScience/EasyHybrid.jl/pull/299). diff --git a/src/config/TrainingConfig.jl b/src/config/TrainingConfig.jl index a6a85447..7a48dd4f 100644 --- a/src/config/TrainingConfig.jl +++ b/src/config/TrainingConfig.jl @@ -22,8 +22,8 @@ $(TYPEDFIELDS) `Optimisers.jl`) three forms are accepted: 1. A single `Optimisers.AbstractRule` (default), applied to the whole - parameter tree, e.g. `Adam(0.01)`. `ps` is wrapped in a - `ComponentArray`. + parameter tree, e.g. `Adam(0.01)`. `ps` is a `NamedTuple` unless + `autodiff_backend` is `AutoForwardDiff`. 2. A `NamedTuple` of rules — one per top-level branch of the parameter tree — e.g. for an `RbQ10` hybrid model: diff --git a/src/losses/compute_loss.jl b/src/losses/compute_loss.jl index 11229c8e..1c6f1ab7 100644 --- a/src/losses/compute_loss.jl +++ b/src/losses/compute_loss.jl @@ -17,11 +17,23 @@ Main loss function for hybrid models that handles both training and evaluation m - In evaluation mode (`logging.train_mode = false`): - `(loss_values, st, ŷ)`: NamedTuple of losses, state and predictions """ +function compute_loss(HM::LuxCore.AbstractLuxContainerLayer, ps, st, data; logging::LoggingLoss) + return compute_loss(HM, ps, st, data, logging) +end + +function compute_loss( + HM::HybridModel{<:Any, <:Vector}, ps, st, + ((x, forcings), (y_t, y_nan)), + logging::LoggingLoss{SymbolicLoss{S}, ExtraLoss{Nothing}, A, true} + ) where {S, A} + loss, st_nn = _hybrid_loss(HM, (x, forcings), ps, st, y_t, y_nan, Val(S), logging.agg) + return loss, ChainRulesCore.ignore_derivatives((; st_nn, fixed = st.fixed)), NamedTuple() +end + function compute_loss( - HM::LuxCore.AbstractLuxContainerLayer, ps, st, ((x, forcings), (y_t, y_nan)); + HM::LuxCore.AbstractLuxContainerLayer, ps, st, ((x, forcings), (y_t, y_nan)), logging::LoggingLoss ) - targets = HM.targets ext_loss = extra_loss(logging) if logging.train_mode diff --git a/src/losses/compute_loss_types.jl b/src/losses/compute_loss_types.jl index 83208663..6f42710e 100644 --- a/src/losses/compute_loss_types.jl +++ b/src/losses/compute_loss_types.jl @@ -3,9 +3,10 @@ export loss_types, training_loss, extra_loss abstract type LossSpec end -struct SymbolicLoss <: LossSpec +struct SymbolicLoss{S} <: LossSpec name::Symbol end +SymbolicLoss(s::Symbol) = SymbolicLoss{s}(s) struct FunctionLoss <: LossSpec f::Function @@ -21,8 +22,8 @@ ParameterizedLoss(f::Function) = ParameterizedLoss(f, (), NamedTuple()) ParameterizedLoss(f::Function, args::Tuple) = ParameterizedLoss(f, args, NamedTuple()) ParameterizedLoss(f::Function, kwargs::NamedTuple) = ParameterizedLoss(f, (), kwargs) -struct ExtraLoss <: LossSpec - f::Union{Function, Nothing} +struct ExtraLoss{F} <: LossSpec + f::F end """ @@ -105,10 +106,10 @@ logging = LoggingLoss( ) ``` """ -struct LoggingLoss{L <: Union{LossSpec, PerTarget}, T <: Function} +struct LoggingLoss{L <: Union{LossSpec, PerTarget}, E <: LossSpec, T <: Function, TM} loss_types::Vector{LossSpec} training_loss::L - extra_loss::LossSpec + extra_loss::E agg::T train_mode::Bool end @@ -125,11 +126,11 @@ function LoggingLoss(; tl = _to_loss_spec(training_loss) el = _to_extra_loss_spec(extra_loss) - return LoggingLoss{typeof(tl), F}(lt, tl, el, agg, train_mode) + return LoggingLoss{typeof(tl), typeof(el), F, train_mode}(lt, tl, el, agg, train_mode) end -_to_loss_spec(s::Symbol) = SymbolicLoss(s) +_to_loss_spec(s::Symbol) = SymbolicLoss{s}(s) _to_loss_spec(f::Function) = _accepts_params(f) ? ParamLoss(f) : FunctionLoss(f) _to_loss_spec(ls::LossSpec) = ls diff --git a/src/models/GenericHybridModel.jl b/src/models/GenericHybridModel.jl index 220254c8..1a8c8d1a 100644 --- a/src/models/GenericHybridModel.jl +++ b/src/models/GenericHybridModel.jl @@ -33,7 +33,7 @@ It combines predictive neural networks (`NNs`) with a `mechanistic_model` to for $(TYPEDFIELDS) """ -struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} +struct HybridModel{T, P, MM, NP, GP, FP, KW, SN, TG} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} "Neural network(s) used to predict parameters. Can be a single `Chain` or a `NamedTuple` of `Chain`s." NNs::T @@ -47,7 +47,7 @@ struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} targets::Vector{Symbol} "The core process-based or mechanistic model function." - mechanistic_model::Function + mechanistic_model::MM "Base parameters of the model (encapsulated in a `ParameterContainer`)." parameters::ParameterContainer @@ -71,6 +71,26 @@ struct HybridModel{T, P} <: LuxCore.AbstractLuxContainerLayer{(:NNs,)} config::NamedTuple end +function HybridModel( + NNs, predictors, forcing, targets, mechanistic_model, parameters, + neural_param_names, global_param_names, fixed_param_names, + scale_nn_outputs, start_from_default, config, + ) + NP, GP, FP = Tuple(neural_param_names), Tuple(global_param_names), Tuple(fixed_param_names) + KW = _accepted_kwarg_names( + mechanistic_model, + Tuple(unique([forcing; neural_param_names; global_param_names; fixed_param_names])), + ) + return HybridModel{ + typeof(NNs), typeof(predictors), typeof(mechanistic_model), + NP, GP, FP, KW, scale_nn_outputs, Tuple(targets), + }( + NNs, predictors, forcing, targets, mechanistic_model, parameters, + neural_param_names, global_param_names, fixed_param_names, + scale_nn_outputs, start_from_default, config, + ) +end + """ _warn_forcing_param_overlap(forcing, param_names) @@ -455,6 +475,75 @@ function _run_nn(m::HybridModel{<:Any, <:Vector}, ds_k::Tuple, ps, st) return scaled_nn_params, (; st_nn = st_nn), (;) end +@generated function _hybrid_loss( + m::HybridModel{<:Any, <:Vector, MM, NP, GP, FP, KW, SN, TG}, + ds_k::Tuple, ps, st, y, y_nan, ::Val{S}, agg + ) where {MM, NP, GP, FP, KW, SN, TG, S} + stmts = Expr[] + names = (NP..., GP..., FP...) + if NP === () + push!(stmts, :(st_nn = st.st_nn)) + else + push!(stmts, :((nn_out, st_nn) = LuxCore.apply(m.NNs, ds_k[1], ps.ps, st.st_nn))) + push!(stmts, :(slices = eachslice(nn_out, dims = 1))) + end + for (i, n) in enumerate(NP) + rhs = SN ? + :(scale_single_param(Val($(QuoteNode(n))), slices[$i], m.parameters)) : + :(slices[$i]) + push!(stmts, Expr(:(=), n, rhs)) + end + for n in GP + push!(stmts, Expr(:(=), n, :(scale_single_param(Val($(QuoteNode(n))), getproperty(ps, $(QuoteNode(n))), m.parameters)))) + end + for n in FP + push!(stmts, Expr(:(=), n, :(getproperty(st.fixed, $(QuoteNode(n)))))) + end + NP === () && GP === () && push!(stmts, :(ChainRulesCore.ignore_derivatives(ps))) + NP === () && KW === () && push!(stmts, :(ChainRulesCore.ignore_derivatives(ds_k))) + !SN && GP === () && push!(stmts, :(ChainRulesCore.ignore_derivatives(m.parameters))) + if KW === nothing + push!(stmts, :(all_params = (; $(names...)))) + push!(stmts, :(y_pred = m.mechanistic_model(; merge(ds_k[2], all_params)...))) + else + args = Expr[] + for n in KW + if n in names + push!(args, Expr(:kw, n, n)) + else + push!(args, Expr(:kw, n, :(getfield(ds_k[2], $(QuoteNode(n)))))) + end + end + push!(stmts, :(y_pred = m.mechanistic_model(; $(args...)))) + end + if TG === () + push!(stmts, :(ChainRulesCore.ignore_derivatives((y, y_nan, agg)))) + push!(stmts, :(return agg(()), st_nn)) + else + terms = Expr[] + for t in TG + qt = QuoteNode(t) + yi = y <: NamedTuple ? :(getfield(y, $qt)) : :(_get_target_y(y, $qt)) + nani = y_nan <: NamedTuple ? :(getfield(y_nan, $qt)) : :(_get_target_y(y_nan, $qt)) + loss = S === :mse ? + :(mean(abs2, ŷ_i[nan_i] .- y_i[nan_i])) : + :(_apply_loss(ŷ_i, y_i, nan_i, $(QuoteNode(S)))) + push!( + terms, :( + let + y_i = $yi + ŷ_i = _get_target_ŷ(y_pred, y_i, $qt) + nan_i = $nani + $loss + end + ) + ) + end + push!(stmts, :(return agg(($(terms...),)), st_nn)) + end + return Expr(:block, stmts...) +end + """ (m::HybridModel)(ds_k::Tuple, ps, st) @@ -510,7 +599,7 @@ function (m::HybridModel)(ds_k::Tuple, ps, st) out = (; y_pred..., extra_params..., parameters = all_params, out_extra...) st_new = (; st_new_nns..., fixed = st.fixed) - return out, st_new + return out, ChainRulesCore.ignore_derivatives(st_new) end function (m::HybridModel)(ds_k, ps, st) diff --git a/src/models/utils.jl b/src/models/utils.jl index 6a17f49e..7c5e69aa 100644 --- a/src/models/utils.jl +++ b/src/models/utils.jl @@ -76,10 +76,18 @@ scaletype(hm::ParameterContainer, name) = hm.scales[name] Map an unconstrained `raw_val` into the parameter's bounds `[ℓ, u]`, using the warp selected for `name` (`:linear`, `:log`, or `:logit`). """ -function scale_single_param(name, raw_val, hm::ParameterContainer) - return _scale(Val(scaletype(hm, name)), lower(hm)[name], upper(hm)[name], raw_val) +function scale_single_param(::Val{N}, raw_val, hm::ParameterContainer) where {N} + bounds = getfield(hm.values, N) + ℓ = bounds[2] + u = bounds[3] + return ℓ .+ (u .- ℓ) .* sigmoid.(raw_val) end +scale_single_param(name::Symbol, raw_val, hm::ParameterContainer) = + scale_single_param(Val(name), raw_val, hm) + +inv_sigmoid(y) = log.(y ./ (1 .- y)) + """ scale_single_param_minmax(name, hm::ParameterContainer) diff --git a/src/training/epoch.jl b/src/training/epoch.jl index a936b845..3f1de436 100644 --- a/src/training/epoch.jl +++ b/src/training/epoch.jl @@ -51,7 +51,7 @@ function build_loss_fn(model, cfg::TrainConfig) extra_loss = cfg.extra_loss, agg = cfg.agg ) - return (model, ps, st, (x, y)) -> compute_loss(model, ps, st, (x, y); logging = logging) + return (model, ps, st, (x, y)) -> compute_loss(model, ps, st, (x, y), logging) end function evaluate_epoch(model, x_train, forcings_train, y_train, mask_train, x_val, forcings_val, y_val, mask_val, ps, st, epoch::Int, init::EpochSnapshot, cfg::TrainConfig) diff --git a/src/training/initialization.jl b/src/training/initialization.jl index 8f5db338..269a3b6b 100644 --- a/src/training/initialization.jl +++ b/src/training/initialization.jl @@ -44,7 +44,7 @@ function init_model_state(model, cfg::TrainConfig) model, ps, st, cfg.opt, opt_state, 0, ) elseif is_optimisers_rule(cfg.opt) - ps = ps |> ComponentArray + cfg.autodiff_backend isa Lux.AutoForwardDiff && (ps = ps |> ComponentArray) Lux.Training.TrainState(model, ps, st, cfg.opt) else ps = ps |> ComponentArray diff --git a/src/training/train_optimization.jl b/src/training/train_optimization.jl index 3c0a3e1a..4ea94d88 100644 --- a/src/training/train_optimization.jl +++ b/src/training/train_optimization.jl @@ -127,7 +127,7 @@ function _build_optim_loss(model, st, cfg::TrainConfig) agg = cfg.agg, ) return function (p, data) - loss_val, _, _ = compute_loss(model, p, st, data; logging) + loss_val, _, _ = compute_loss(model, p, st, data, logging) return loss_val end end diff --git a/src/utils/ad_rules.jl b/src/utils/ad_rules.jl new file mode 100644 index 00000000..d81e2c18 --- /dev/null +++ b/src/utils/ad_rules.jl @@ -0,0 +1,6 @@ +# Zygote traces Base's keyword-arg name sorting; the results are Bool/Symbol. +for (fname, nargs) in ((:diff_names, 2), (:sym_in, 2), (:merge_names, 2)) + isdefined(Base, fname) || continue + args = ntuple(_ -> :(::Any), nargs) + @eval ChainRulesCore.@non_differentiable Base.$fname($(args...)) +end diff --git a/src/utils/utils.jl b/src/utils/utils.jl index 5d84e255..405d3ff2 100644 --- a/src/utils/utils.jl +++ b/src/utils/utils.jl @@ -1,3 +1,4 @@ +include("ad_rules.jl") include("macro_hybrid.jl") include("wrap_tuples.jl") include("tools.jl") diff --git a/test/runtests.jl b/test/runtests.jl index d5735db4..b1294986 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -17,4 +17,5 @@ include("test_wrap_tuples.jl") include("test_extract_weights.jl") include("test_transformers.jl") include("test_seq2seq.jl") +include("test_synth_performance.jl") include("test_zarr_io.jl") diff --git a/test/test_compute_loss.jl b/test/test_compute_loss.jl index 42bb57ff..a67e8d85 100644 --- a/test/test_compute_loss.jl +++ b/test/test_compute_loss.jl @@ -5,6 +5,7 @@ using Statistics using DimensionalData using Random using DataFrames +using Zygote @testset "_compute_loss" begin # Test data setup @@ -305,6 +306,16 @@ end ŷ_actual, y_t, y_nan, targets, :mse, sum ) @test loss_value ≈ main_loss + + g_fused = only(Zygote.gradient(p -> compute_loss(HM, p, st, (data[1], (data[2], y_nan)), logging)[1], ps)) + g_pub = only( + Zygote.gradient( + p -> _compute_loss(HM(data[1], p, st)[1], y_t, y_nan, Tuple(targets), :mse, sum), + ps, + ) + ) + @test g_fused.b ≈ g_pub.b + @test g_fused.ps.layer_2.weight ≈ g_pub.ps.layer_2.weight end @testset "Evaluation mode with extra_loss" begin diff --git a/test/test_generic_hybrid_model.jl b/test/test_generic_hybrid_model.jl index 228a2ca8..76570ee3 100644 --- a/test/test_generic_hybrid_model.jl +++ b/test/test_generic_hybrid_model.jl @@ -2,6 +2,7 @@ using Lux using Random using AxisKeys using ComponentArrays +using Zygote # Test data generation dk = gen_linear_data() @@ -255,6 +256,11 @@ end @test haskey(output.parameters, :c) @test haskey(new_st, :st_nn) @test haskey(new_st, :fixed) + + # Gradient of a prediction-only loss (returned state unused) + g = Zygote.gradient(p -> sum(abs2, model(data_[1], p, st)[1].y_pred), ps) + @test !isnothing(g[1]) + @test g[1].b isa AbstractVector end @testset "HybridModel with scale_nn_outputs=true" begin diff --git a/test/test_synth_performance.jl b/test/test_synth_performance.jl new file mode 100644 index 00000000..a91101c3 --- /dev/null +++ b/test/test_synth_performance.jl @@ -0,0 +1,166 @@ +using Random +using Statistics +using Zygote +using Lux +using Logging +using DataFrames +using EasyHybrid +using EasyHybrid: compute_loss, scale_single_param, LoggingLoss + +# Synthetic RbQ10 setup matching docs/literate/research/synthetic_respiration.jl +# (same columns, hidden layers, activation, scaling, and batch size; data generated locally). + +const SYNTH_BATCH = 512 +# Paired median of fused `compute_loss` Zygote.gradient vs a Lux same-math gradient +# on this example. New code may not exceed this by more than `SYNTH_SLOWDOWN_TOLERANCE`. +const SYNTH_COMPUTE_LOSS_VS_LUX = 1.18 +const SYNTH_SLOWDOWN_TOLERANCE = 0.1 + +function _synth_rbq10(; n = 2048, batch = SYNTH_BATCH, seed = 42) + rng = MersenneTwister(seed) + ta = 10 .+ 10 .* randn(rng, n) + sw_pot = abs.(50 .+ 20 .* randn(rng, n)) + dsw_pot = vcat(0.0, diff(sw_pot)) + true_Q10 = 2.0 + true_rb = 3.0 .+ 0.02 .* (sw_pot .- mean(sw_pot)) + reco = true_rb .* (true_Q10 .^ (0.1 .* (ta .- 15.0))) .+ 0.1 .* randn(rng, n) + df = DataFrame(; + ta = Float32.(ta), + sw_pot = Float32.(sw_pot), + dsw_pot = Float32.(dsw_pot), + reco = Float32.(reco), + ) + + function RbQ10(; ta, Q10, rb, tref = 15.0f0) + reco = rb .* Q10 .^ (0.1f0 .* (ta .- tref)) + return (; reco, Q10, rb) + end + parameters = ( + rb = (3.0f0, 0.0f0, 13.0f0), + Q10 = (2.0f0, 1.0f0, 4.0f0), + ) + hm = constructHybridModel( + [:sw_pot, :dsw_pot], [:ta], [:reco], RbQ10, parameters, [:rb], [:Q10]; + hidden_layers = [16, 16], + activation = sigmoid, + scale_nn_outputs = true, + input_batchnorm = true, + ) + (x_all, forcings_all), y_all = prepare_data(hm, df) + x = Array(x_all)[:, 1:batch] + ta_b = Array(forcings_all.ta)[1:batch] + yv = Array(y_all.reco)[1:batch] + mask = .!isnan.(yv) + ds = ((x, (; ta = ta_b)), ((; reco = yv), (; reco = mask))) + ps, st = LuxCore.setup(MersenneTwister(seed), hm) + logging = LoggingLoss(; training_loss = :mse, loss_types = [:mse], train_mode = true, agg = sum) + return (; hm, ps, st, ds, logging, df, x, ta_b, yv, mask) +end + +function _median_s(f; warmup = 4, samples = 21) + for _ in 1:warmup + f() + end + GC.gc() + times = Vector{Float64}(undef, samples) + for i in 1:samples + times[i] = @elapsed f() + end + return median(times) +end + +function _paired_median_ratio(f_a, f_b; warmup = 8, samples = 41) + for _ in 1:warmup + f_a() + f_b() + end + GC.gc() + ratios = Vector{Float64}(undef, samples) + for i in 1:samples + if isodd(i) + ta = @elapsed f_a() + tb = @elapsed f_b() + else + tb = @elapsed f_b() + ta = @elapsed f_a() + end + ratios[i] = ta / tb + end + return median(ratios) +end + +function _lux_same_math(s) + hm, st, x, ta, yv, mask = s.hm, s.st, s.x, s.ta_b, s.yv, s.mask + bounds = hm.parameters + return function (p) + o, _ = LuxCore.apply(hm.NNs, x, p.ps, st.st_nn) + rb = scale_single_param(:rb, eachslice(o, dims = 1)[1], bounds) + Q10 = scale_single_param(:Q10, p.Q10, bounds) + reco = rb .* Q10 .^ (0.1f0 .* (ta .- 15.0f0)) + return mean(abs2, reco[mask] .- yv[mask]) + end +end + +@testset "Synthetic RbQ10 performance" begin + with_logger(NullLogger()) do + s = _synth_rbq10() + hm, ps, st, ds, logging = s.hm, s.ps, s.st, s.ds, s.logging + lux_loss = _lux_same_math(s) + loss_fn = (model, p, st_, data) -> compute_loss(model, p, st_, data, logging) + + f_fwd = () -> hm((s.x, (; ta = s.ta_b)), ps, st) + f_loss = () -> compute_loss(hm, ps, st, ds, logging) + f_lux = () -> Zygote.gradient(lux_loss, ps) + f_cl = () -> Zygote.gradient(p -> compute_loss(hm, p, st, ds, logging)[1], ps) + f_hm_mse = () -> Zygote.gradient( + p -> begin + out, _ = hm((s.x, (; ta = s.ta_b)), p, st) + mean(abs2, out.reco[s.mask] .- s.yv[s.mask]) + end, ps + ) + f_step = () -> begin + ts = Lux.Training.TrainState(hm, ps, st, AdamW(0.1)) + Lux.Training.single_train_step!(AutoZygote(), loss_fn, ds, ts) + end + + f_lux(); f_fwd(); f_cl(); f_hm_mse(); f_step() + + t_fwd = _median_s(f_fwd) + t_loss = _median_s(f_loss) + t_lux = _median_s(f_lux) + t_cl = _median_s(f_cl) + t_hm_mse = _median_s(f_hm_mse) + t_step = _median_s(f_step) + ratio = _paired_median_ratio(f_cl, f_lux) + cap = SYNTH_COMPUTE_LOSS_VS_LUX * (1 + SYNTH_SLOWDOWN_TOLERANCE) + + println("synthetic RbQ10 steps (median µs):") + println(" HybridModel forward ", round(t_fwd * 1.0e6; digits = 1)) + println(" compute_loss value ", round(t_loss * 1.0e6; digits = 1)) + println(" Lux same-math gradient ", round(t_lux * 1.0e6; digits = 1)) + println(" compute_loss gradient ", round(t_cl * 1.0e6; digits = 1)) + println(" HybridModel + MSE gradient ", round(t_hm_mse * 1.0e6; digits = 1)) + println(" single_train_step! ", round(t_step * 1.0e6; digits = 1)) + println(" compute_loss / Lux (paired) ", round(ratio; digits = 3), " (cap ", round(cap; digits = 3), ")") + + @testset "step: HybridModel forward" begin + @test t_fwd > 0 + end + @testset "step: compute_loss value" begin + @test t_loss > 0 + cl = compute_loss(hm, ps, st, ds, logging)[1] + out, _ = hm((s.x, (; ta = s.ta_b)), ps, st) + @test cl ≈ mean(abs2, out.reco[s.mask] .- s.yv[s.mask]) rtol = 1.0e-4 + @test cl ≈ lux_loss(ps) rtol = 1.0e-4 + end + @testset "step: compute_loss gradient vs Lux" begin + @test ratio ≤ SYNTH_COMPUTE_LOSS_VS_LUX * (1 + SYNTH_SLOWDOWN_TOLERANCE) + end + @testset "step: compute_loss not slower than public HybridModel MSE" begin + @test t_cl ≤ t_hm_mse * (1 + SYNTH_SLOWDOWN_TOLERANCE) + end + @testset "step: single_train_step!" begin + @test t_step > 0 + end + end +end