Fix Optimisers.adjust! to mutate Training.TrainState.optimizer

[Copilot]

Optimisers.adjust!(::Training.TrainState, ...) updated optimizer_state but did not actually mutate the optimizer field on the passed TrainState. As a result, train_state.optimizer could remain stale after in-place adjust calls.

• Behavior fix in Training helper API

  • Updated Optimisers.adjust!(ts::TrainState, eta::Real) and Optimisers.adjust!(ts::TrainState; kwargs...) to directly assign ts.optimizer = Optimisers.adjust(...).
  • This preserves in-place semantics for adjust! on TrainState.

• Regression coverage for optimizer consistency

  • Extended training API tests to assert both:
    • tstate.optimizer_state...rule.eta
    • tstate.optimizer.eta
  • Added assertions for both adjust and adjust! overloads (eta positional and keyword forms).

function Optimisers.adjust!(ts::TrainState, eta::Real)
    Optimisers.adjust!(ts.optimizer_state, eta)
    ts.optimizer = Optimisers.adjust(ts.optimizer, eta)
    return ts
end

[github-actions]

Benchmark Results (Julia v1.12)

Time benchmarks

	main	f76fd73...	main / f76fd73...
basics/MHA	3.8 ± 0.42 μs	3.81 ± 0.42 μs	0.995 ± 0.16
basics/MHA (first run)	4.08 ± 0.51 μs	4.14 ± 0.57 μs	0.985 ± 0.18
basics/MHA reactant	0.198 ± 0.28 ms	0.205 ± 0.25 ms	0.969 ± 1.8
basics/MHA reactant (comp + run)	0.197 ± 0.008 s	0.198 ± 0.0066 s	0.997 ± 0.052
basics/conv	12.5 ± 8.4 μs	12.6 ± 8.5 μs	0.995 ± 0.95
basics/conv (first run)	12.6 ± 2.5 μs	12.6 ± 2.3 μs	1 ± 0.27
basics/conv reactant	0.0857 ± 0.011 ms	0.0829 ± 0.011 ms	1.03 ± 0.19
basics/conv reactant (comp + run)	0.163 ± 0.0021 s	0.167 ± 0.0085 s	0.978 ± 0.052
basics/dense	0.171 ± 0.031 μs	0.17 ± 0.02 μs	1.01 ± 0.22
basics/dense (first run)	0.211 ± 0.011 μs	0.23 ± 0.02 μs	0.917 ± 0.093
basics/dense reactant	0.0785 ± 0.0033 ms	0.0821 ± 0.011 ms	0.955 ± 0.13
basics/dense reactant (comp + run)	0.127 ± 0.0044 s	0.128 ± 0.01 s	0.987 ± 0.088
time_to_load	0.65 ± 0.0038 s	0.643 ± 0.012 s	1.01 ± 0.019

Memory benchmarks

	main	f76fd73...	main / f76fd73...
basics/MHA	0.087 k allocs: 5.81 kB	0.085 k allocs: 5.75 kB	1.01
basics/MHA (first run)	0.087 k allocs: 5.81 kB	0.085 k allocs: 5.75 kB	1.01
basics/MHA reactant	19 allocs: 0.578 kB	19 allocs: 0.578 kB	1
basics/MHA reactant (comp + run)	19.6 k allocs: 1.77 MB	19.6 k allocs: 1.77 MB	1
basics/conv	0.039 k allocs: 4.55 kB	0.039 k allocs: 4.55 kB	1
basics/conv (first run)	0.039 k allocs: 4.55 kB	0.039 k allocs: 4.55 kB	1
basics/conv reactant	15 allocs: 0.438 kB	15 allocs: 0.438 kB	1
basics/conv reactant (comp + run)	7.61 k allocs: 1.24 MB	7.61 k allocs: 1.24 MB	1
basics/dense	2 allocs: 0.109 kB	2 allocs: 0.109 kB	1
basics/dense (first run)	2 allocs: 0.109 kB	2 allocs: 0.109 kB	1
basics/dense reactant	15 allocs: 0.422 kB	15 allocs: 0.422 kB	1
basics/dense reactant (comp + run)	4.92 k allocs: 1.13 MB	4.92 k allocs: 1.13 MB	1
time_to_load	0.145 k allocs: 11 kB	0.145 k allocs: 11 kB	1