Pass train_state to update_params

[Niccolo-Ajroldi]

This pull-request adds train_state to the arguments of update_params.

This allows for example to implement a time-based LR scheduler (see #785).

Since train_state is a dictionary, we have to avoid that a submission modifies its content.
The solution adopted here passes train_state by copy. We use a shallow copying, since all values in train_state are primitive types, and cannot be modified in-place. Although very fast, this copying operation would however be counted as part of the submission time.

A possible alternative solution that would avoid this problem would be to pass train_state['accumulated_submission_time'] directly instead of passing the entire dictionary.

[github-actions]

MLCommons CLA bot All contributors have signed the MLCommons CLA ✍️ ✅

[priyakasimbeg]

lgtm