# How to correctly use Lux and Training API with reactant?
**URL:** https://discourse.julialang.org/t/how-to-correctly-use-lux-and-training-api-with-reactant/133799
**Category:** Machine Learning
**Tags:** lux
**Created:** [November 11, 2025, 5:23pm UTC](https://discourse.julialang.org/t/how-to-correctly-use-lux-and-training-api-with-reactant/133799 "2025-11-11T17:23:55Z")
**Posts on this page:** 4
**Page:** 1
### Author:  [@rkube](https://discourse.julialang.org/u/rkube)
#### Post date: [November 11, 2025, 5:23pm UTC](https://discourse.julialang.org/t/how-to-correctly-use-lux-and-training-api-with-reactant/133799/1 "2025-11-11T17:23:55Z")
Hi,
I’m trying to train neural networks with Lux, Reactant, and the Training API and am confused about how to tie it together.
1. When I’m compiling a model with reactant, should I still use the TrainState API?
I tried passing a compiled model into TrainState and get this error:
```julia-auto
julia> tstate = Training.TrainState(model_compiled, ps_ra, st_ra, opt)
ERROR: MethodError: no method matching Lux.Training.TrainState(::Reactant.Compiler.Thunk{…}, ::@NamedTuple{…}, ::@NamedTuple{…}, ::Adam{…})
The type `Lux.Training.TrainState` exists, but no method is defined for this combination of argument types when trying to construct it.
Closest candidates are:
Lux.Training.TrainState(:: __T_cache, ::__ T_objective_function, :: __T_model, ::__ T_parameters, :: __T_states, ::__ T_optimizer, ::__T_optimizer_state, ::Int64) where {__T_cache, __T_objective_function,__ T_model, __T_parameters,__ T_states, __T_optimizer,__ T_optimizer_state}
@ Lux ~/.julia/packages/ConcreteStructs/7Lv7u/src/ConcreteStructs.jl:142
Lux.Training.TrainState(::AbstractLuxLayer, ::Any, ::Any, ::AbstractRule)
@ Lux ~/.julia/packages/Lux/uEbqO/src/helpers/training.jl:122
Stacktrace:
[1] top-level scope
@ REPL[48]:1
Some type information was truncated. Use `show(err)` to see complete types.
```
1. The docs for [TrainState](https://lux.csail.mit.edu/stable/api/Lux/utilities#Training-API) say that the version is not stable and I should use the version with the Optimisers.jl API. Where can I find this API? The tutorials, f.ex. the MLP fitting, still refer to this one.
---
### Author:  [@Sebastian\_Schmitt](https://discourse.julialang.org/u/Sebastian_Schmitt)
#### Post date: [November 11, 2025, 5:53pm UTC](https://discourse.julialang.org/t/how-to-correctly-use-lux-and-training-api-with-reactant/133799/2 "2025-11-11T17:53:30Z")
Hi, you need to use the default Lux model (not compiled) for constructing the train state. The model will be compiled (and cached) when you call `single_train_step` (or `compute_gradients`) the first time. The hint on in the TrainState docs is referring to creating the TrainState by its default (struct) constructor, i.e. `TrainState(fields…)` (at least that is how I undertand it). The API you’re using (`TrainState(model, ps, st, opt)`) is the Optimiser.jl API already.
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [November 12, 2025, 4:57pm UTC](https://discourse.julialang.org/t/how-to-correctly-use-lux-and-training-api-with-reactant/133799/3 "2025-11-12T16:57:07Z")
@Sebastian_Schmitt’s response is the correct solution.
This seems to be a common mistake and from the next release of Lux, users will see the following error when attempting to pass a compiled function in TrainState
````julia-auto
julia> train_state = Training.TrainState(model_compiled, ps, st, Adam())
ERROR: ArgumentError: Invalid TrainState construction using a compiled function.
`TrainState` is being constructed with a reactant compiled function, i.e. a
`Reactant.Compiler.Thunk`. This is likely a mistake as the model should be
passed in directly without being compiled first.
This is likely originating from the following style of usage:
```julia
using Lux, Reactant, Random, Optimisers
rdev = reactant_device()
model = Dense(10, 10)
ps, st = Lux.setup(Random.default_rng(), model) |> rdev
x = rand(10) |> rdev
model_compiled = @compile model(x, ps, st)
train_state = Training.TrainState(model_compiled, ps, st, Adam())
```
Instead avoid compiling the model and pass it directly to `TrainState`. When
`single_train_step` or other function are called on the `TrainState`, the
model will be compiled automatically.
```julia
train_state = Training.TrainState(model, ps, st, Adam())
```
For end-to-end usage example refer to the documentation:
Stacktrace:
[1] Lux.Training.TrainState(::Reactant.Compiler.Thunk{…}, ps::@NamedTuple{…}, st::@NamedTuple{}, optimizer::Adam{…})
@ LuxReactantExt /mnt/software/lux/Lux.jl/ext/LuxReactantExt/training.jl:5
[2] top-level scope
@ REPL[42]:1
[3] top-level scope
@ none:1
Some type information was truncated. Use `show(err)` to see complete types.
````
---
### Author:  [@rkube](https://discourse.julialang.org/u/rkube)
#### Post date: [November 12, 2025, 6:04pm UTC](https://discourse.julialang.org/t/how-to-correctly-use-lux-and-training-api-with-reactant/133799/4 "2025-11-12T18:04:09Z")
Awesome, thank you. A suggestion on the error message - put the correct code on top. Something like
````julia-auto
When passing a model into TrainState make sure it is not compiled. When
`single_train_step` or other function are called on the `TrainState`, the
model will be compiled automatically.
```julia
using Lux, Reactant, Random, Optimisers
rdev = reactant_device()
model = Dense(10, 10)
ps, st = Lux.setup(Random.default_rng(), model) |> rdev
x = rand(10) |> rdev
train_state = Training.TrainState(model, ps, st, Adam())
````
Do _not_ pass a compiled model into Training.TrainState
```julia
model_compiled = @compile model(x, ps, st)
train_state = Training.TrainState(model_compiled, ps, st, Adam())
-> This throws an error
```
```julia-auto
```