def federated_sum(x):
py_typecheck.check_type(x, building_blocks.ComputationBuildingBlock)
operand_type = x.type_signature.member
zero = building_block_factory.create_generic_constant(operand_type, 0)
plus_op = building_block_factory.create_tensorflow_binary_operator_with_upcast(
tf.add, computation_types.StructType([operand_type, operand_type]))
identity = building_block_factory.create_identity(operand_type)
return building_block_factory.create_federated_aggregate(
x, zero, plus_op, plus_op, identity)
def federated_sum(x):
zero = value_impl.ValueImpl(
building_block_factory.create_generic_constant(x.type_signature.member,
0), context_stack)
plus_op = value_impl.ValueImpl(
building_block_factory.create_tensorflow_binary_operator_with_upcast(
computation_types.StructType(
[x.type_signature.member, x.type_signature.member]), tf.add),
context_stack)
return federated_reduce([x, zero, plus_op])
async def _num_clients(executor):
"""Returns the number of clients for the given `executor`."""
# We implement num_clients as a federated_aggregate to allow for federated
# op resolving strategies which implement only the minimal set of
# intrinsics.
int_at_clients_type = computation_types.at_clients(tf.int32)
zero_type = tf.int32
accumulate_type = computation_types.FunctionType(
computation_types.StructType([tf.int32, tf.int32]), tf.int32)
merge_type = accumulate_type
report_type = computation_types.FunctionType(tf.int32, tf.int32)
intrinsic_type = computation_types.FunctionType(
computation_types.StructType([
int_at_clients_type, zero_type, accumulate_type,
merge_type, report_type
]), computation_types.at_server(tf.int32))
intrinsic = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_AGGREGATE, intrinsic_type)
add_comp = building_block_factory.create_tensorflow_binary_operator_with_upcast(
computation_types.StructType([tf.int32, tf.int32]),
tf.add).proto
identity_comp = building_block_factory.create_compiled_identity(
computation_types.TensorType(tf.int32)).proto
fn, client_data, zero_value, add_value, identity_value = await asyncio.gather(
executor.create_value(intrinsic, intrinsic_type),
executor.create_value(
1, computation_types.at_clients(tf.int32, all_equal=True)),
executor.create_value(0, tf.int32),
executor.create_value(add_comp),
executor.create_value(identity_comp))
arg = await executor.create_struct([
client_data, zero_value, add_value, add_value, identity_value
])
call = await executor.create_call(fn, arg)
result = await call.compute()
if isinstance(result, tf.Tensor):
return result.numpy()
else:
return result
async def compute_intrinsic_federated_weighted_mean(
executor: executor_base.Executor, arg: executor_value_base.ExecutorValue
) -> executor_value_base.ExecutorValue:
"""Computes a federated weighted mean on the given `executor`.
Args:
executor: The executor to use.
arg: The argument to embedded in `executor`.
Returns:
The result embedded in `executor`.
"""
type_analysis.check_valid_federated_weighted_mean_argument_tuple_type(
arg.type_signature)
zip1_type = computation_types.FunctionType(
computation_types.StructType([
computation_types.at_clients(arg.type_signature[0].member),
computation_types.at_clients(arg.type_signature[1].member)
]),
computation_types.at_clients(
computation_types.StructType(
[arg.type_signature[0].member, arg.type_signature[1].member])))
multiply_blk = building_block_factory.create_tensorflow_binary_operator_with_upcast(
zip1_type.result.member, tf.multiply)
map_type = computation_types.FunctionType(
computation_types.StructType(
[multiply_blk.type_signature, zip1_type.result]),
computation_types.at_clients(multiply_blk.type_signature.result))
sum1_type = computation_types.FunctionType(
computation_types.at_clients(map_type.result.member),
computation_types.at_server(map_type.result.member))
sum2_type = computation_types.FunctionType(
computation_types.at_clients(arg.type_signature[1].member),
computation_types.at_server(arg.type_signature[1].member))
zip2_type = computation_types.FunctionType(
computation_types.StructType([sum1_type.result, sum2_type.result]),
computation_types.at_server(
computation_types.StructType(
[sum1_type.result.member, sum2_type.result.member])))
divide_blk = building_block_factory.create_tensorflow_binary_operator_with_upcast(
zip2_type.result.member, tf.divide)
async def _compute_multiply_fn():
return await executor.create_value(multiply_blk.proto,
multiply_blk.type_signature)
async def _compute_multiply_arg():
zip1_comp = create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS, zip1_type)
zip_fn = await executor.create_value(zip1_comp, zip1_type)
return await executor.create_call(zip_fn, arg)
async def _compute_product_fn():
map_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_MAP,
map_type)
return await executor.create_value(map_comp, map_type)
async def _compute_product_arg():
multiply_fn, multiply_arg = await asyncio.gather(
_compute_multiply_fn(), _compute_multiply_arg())
return await executor.create_struct((multiply_fn, multiply_arg))
async def _compute_products():
product_fn, product_arg = await asyncio.gather(_compute_product_fn(),
_compute_product_arg())
return await executor.create_call(product_fn, product_arg)
async def _compute_total_weight():
sum2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM,
sum2_type)
sum2_fn, sum2_arg = await asyncio.gather(
executor.create_value(sum2_comp, sum2_type),
executor.create_selection(arg, index=1))
return await executor.create_call(sum2_fn, sum2_arg)
async def _compute_sum_of_products():
sum1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM,
sum1_type)
sum1_fn, products = await asyncio.gather(
executor.create_value(sum1_comp, sum1_type), _compute_products())
return await executor.create_call(sum1_fn, products)
async def _compute_zip2_fn():
zip2_comp = create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_SERVER, zip2_type)
return await executor.create_value(zip2_comp, zip2_type)
async def _compute_zip2_arg():
sum_of_products, total_weight = await asyncio.gather(
_compute_sum_of_products(), _compute_total_weight())
return await executor.create_struct([sum_of_products, total_weight])
async def _compute_divide_fn():
return await executor.create_value(divide_blk.proto,
divide_blk.type_signature)
async def _compute_divide_arg():
zip_fn, zip_arg = await asyncio.gather(_compute_zip2_fn(),
_compute_zip2_arg())
return await executor.create_call(zip_fn, zip_arg)
async def _compute_apply_fn():
apply_type = computation_types.FunctionType(
computation_types.StructType(
[divide_blk.type_signature, zip2_type.result]),
computation_types.at_server(divide_blk.type_signature.result))
apply_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_APPLY,
apply_type)
return await executor.create_value(apply_comp, apply_type)
async def _compute_apply_arg():
divide_fn, divide_arg = await asyncio.gather(_compute_divide_fn(),
_compute_divide_arg())
return await executor.create_struct([divide_fn, divide_arg])
async def _compute_divided():
apply_fn, apply_arg = await asyncio.gather(_compute_apply_fn(),
_compute_apply_arg())
return await executor.create_call(apply_fn, apply_arg)
return await _compute_divided()
async def compute_intrinsic_federated_weighted_mean(
executor: executor_base.Executor,
arg: executor_value_base.ExecutorValue,
local_computation_factory: local_computation_factory_base.
LocalComputationFactory = tensorflow_computation_factory.
TensorFlowComputationFactory()
) -> executor_value_base.ExecutorValue:
"""Computes a federated weighted mean on the given `executor`.
Args:
executor: The executor to use.
arg: The argument to embedded in `executor`.
local_computation_factory: An instance of `LocalComputationFactory` to use
to construct local computations used as parameters in certain federated
operators (such as `tff.federated_sum`, etc.). Defaults to a TensorFlow
computation factory that generates TensorFlow code.
Returns:
The result embedded in `executor`.
"""
type_analysis.check_valid_federated_weighted_mean_argument_tuple_type(
arg.type_signature)
zip1_type = computation_types.FunctionType(
computation_types.StructType([
computation_types.at_clients(arg.type_signature[0].member),
computation_types.at_clients(arg.type_signature[1].member)
]),
computation_types.at_clients(
computation_types.StructType(
[arg.type_signature[0].member, arg.type_signature[1].member])))
operand_type = zip1_type.result.member[0]
scalar_type = zip1_type.result.member[1]
multiply_comp_pb, multiply_comp_type = local_computation_factory.create_scalar_multiply_operator(
operand_type, scalar_type)
multiply_blk = building_blocks.CompiledComputation(
multiply_comp_pb, type_signature=multiply_comp_type)
map_type = computation_types.FunctionType(
computation_types.StructType(
[multiply_blk.type_signature, zip1_type.result]),
computation_types.at_clients(multiply_blk.type_signature.result))
sum1_type = computation_types.FunctionType(
computation_types.at_clients(map_type.result.member),
computation_types.at_server(map_type.result.member))
sum2_type = computation_types.FunctionType(
computation_types.at_clients(arg.type_signature[1].member),
computation_types.at_server(arg.type_signature[1].member))
zip2_type = computation_types.FunctionType(
computation_types.StructType([sum1_type.result, sum2_type.result]),
computation_types.at_server(
computation_types.StructType(
[sum1_type.result.member, sum2_type.result.member])))
divide_blk = building_block_factory.create_tensorflow_binary_operator_with_upcast(
zip2_type.result.member, tf.divide)
async def _compute_multiply_fn():
return await executor.create_value(multiply_blk.proto,
multiply_blk.type_signature)
async def _compute_multiply_arg():
zip1_comp = create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS, zip1_type)
zip_fn = await executor.create_value(zip1_comp, zip1_type)
return await executor.create_call(zip_fn, arg)
async def _compute_product_fn():
map_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_MAP,
map_type)
return await executor.create_value(map_comp, map_type)
async def _compute_product_arg():
multiply_fn, multiply_arg = await asyncio.gather(
_compute_multiply_fn(), _compute_multiply_arg())
return await executor.create_struct((multiply_fn, multiply_arg))
async def _compute_products():
product_fn, product_arg = await asyncio.gather(_compute_product_fn(),
_compute_product_arg())
return await executor.create_call(product_fn, product_arg)
async def _compute_total_weight():
sum2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM,
sum2_type)
sum2_fn, sum2_arg = await asyncio.gather(
executor.create_value(sum2_comp, sum2_type),
executor.create_selection(arg, 1))
return await executor.create_call(sum2_fn, sum2_arg)
async def _compute_sum_of_products():
sum1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM,
sum1_type)
sum1_fn, products = await asyncio.gather(
executor.create_value(sum1_comp, sum1_type), _compute_products())
return await executor.create_call(sum1_fn, products)
async def _compute_zip2_fn():
zip2_comp = create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_SERVER, zip2_type)
return await executor.create_value(zip2_comp, zip2_type)
async def _compute_zip2_arg():
sum_of_products, total_weight = await asyncio.gather(
_compute_sum_of_products(), _compute_total_weight())
return await executor.create_struct([sum_of_products, total_weight])
async def _compute_divide_fn():
return await executor.create_value(divide_blk.proto,
divide_blk.type_signature)
async def _compute_divide_arg():
zip_fn, zip_arg = await asyncio.gather(_compute_zip2_fn(),
_compute_zip2_arg())
return await executor.create_call(zip_fn, zip_arg)
async def _compute_apply_fn():
apply_type = computation_types.FunctionType(
computation_types.StructType(
[divide_blk.type_signature, zip2_type.result]),
computation_types.at_server(divide_blk.type_signature.result))
apply_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_APPLY,
apply_type)
return await executor.create_value(apply_comp, apply_type)
async def _compute_apply_arg():
divide_fn, divide_arg = await asyncio.gather(_compute_divide_fn(),
_compute_divide_arg())
return await executor.create_struct([divide_fn, divide_arg])
async def _compute_divided():
apply_fn, apply_arg = await asyncio.gather(_compute_apply_fn(),
_compute_apply_arg())
return await executor.create_call(apply_fn, apply_arg)
return await _compute_divided()
