Update docstring of `tff.learning.state_with_new_model_weights`.
@@ -126,6 +126,8 @@ class ServerState(object):
model_broadcast_state = attr.ib()
+# TODO(b/227484749): standardize/improve the function for
+# `tff.learning.algorithms`.
def state_with_new_model_weights(
server_state: ServerState,
trainable_weights: List[np.ndarray],
@@ -133,6 +135,9 @@ def state_with_new_model_weights(
) -> ServerState:
"""Returns a `ServerState` with updated model weights.
+ This function cannot currently be used for `tff.learning.algorithms`, which
+ can have a different `ServerState` structure.
+
Args:
server_state: A server state object returned by an iterative training
process like `tff.learning.build_federated_averaging_process`.
