Help understanding non-mutating function in Lux.Training

Specific Domains Machine Learning
1

The Training.apply_gradient non-mutating function is implemented as below in Lux.jl:

function apply_gradients(ts::TrainState, grads)
    optimizer_state, ps = Optimisers.update(ts.optimizer_state, ts.parameters, grads)
    @set! ts.parameters = ps
    @set! ts.optimizer_state = optimizer_state
    @set! ts.step = ts.step + 1
    return ts
end

The way I understand is, @set! modifies the struct in place. But how does it still behave as a non-mutating function (the output of the function is the updated ts, but the ts passed into the function remains unchanged)?

2

See Introduction · Setfield.jl. @set! doesn’t mutate the struct, instead it updates the struct field → creates a new struct and then updates the variable ts to reference the newly created struct.