def local_train(initial_model, learning_rate, all_batches):
@tff.federated_computation(model_type, batch_type)
def batch_fn(model, batch):
return batch_train(model, batch, learning_rate)
return tff.sequence_reduce(all_batches, initial_model, batch_fn)
def local_train(initial_model, learning_rate, all_batches):
@tff.federated_computation(model_type, batch_type)
def batch_fn(model, batch):
return batch_train(model, batch, learning_rate)
# Blocked on implementing sequence_reduce in local executor.
return tff.sequence_reduce(all_batches, initial_model, batch_fn)