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()