secure-aggregation.py

# compute and collect all model gradients as private inputs
model_grads = zip(*[
    tfe.define_private_input(
        data_owner.player_name,
        data_owner.compute_gradient)
    for data_owner in data_owners
])

# compute mean gradient securely
aggregated_model_grads = [
    tfe.add_n(grads) / len(grads)
    for grads in model_grads
]

# reveal only aggregated gradients to model owner
iteration_op = tfe.define_output(
    model_owner.player_name,
    model_owner.update_model,
    aggregated_model_grads)

with tfe.Session() as sess:
    for _ in range(num_iterations):
        sess.run(iteration_op)