Question: how can one use TrainState.cache?

[luisaforozco]

From the TrainState docs it seems that there is a field in the struct called cache that can be used to store anything.
However the constructor sets the cache to nothing disallowing the user to initialize the cache field.

What is the right way to initialize the TrainState.cache?
Concretely, I would like to initialize and/or modify TrainState.cache to a Dict or NamedTuple in which I store the best parameters and loss.
Something like:

state = TrainState(params, opt_state, Dict(:best_params => deepcopy(params), :min_loss => Inf))
# and later in the training loop ...
# Update the minimum loss and best parameters in the cache
if loss < state.cache[:min_loss]
    state.cache[:min_loss] = loss
    state.cache[:best_params] = deepcopy(new_params)
end

[avik-pal]

cache is an internal field. It is used by the backend implementations, i.e. the ones defined in the extensions like Zygote, Enzyme, etc. to cache intermediate buffers and such. It is not meant to be exposed to the external user.

For your case, you should simply create an external cache and update it as necessary without involving the TrainState.