# Update the model with new state
nnx.update(model, new_state)

        __
        __