def test_serialize_deserialize_federated_at_clients(self): x = [10, 20] x_type = type_factory.at_clients(tf.int32) value_proto, value_type = executor_service_utils.serialize_value(x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), '{int32}@CLIENTS') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertEqual(y, [10, 20])
def test_returns_value_with_federated_type_at_clients_all_equal(self): value = [eager_tf_executor.EagerValue(10.0, None, tf.float32)] type_signature = type_factory.at_clients(tf.float32, all_equal=True) value = federated_resolving_strategy.FederatedResolvingStrategyValue( value, type_signature) result = self.run_sync(value.compute()) self.assertEqual(result, 10.0)
async def compute_federated_zip_at_clients( self, arg: FederatedComposingStrategyValue ) -> FederatedComposingStrategyValue: py_typecheck.check_type(arg.type_signature, computation_types.StructType) py_typecheck.check_len(arg.type_signature, 2) py_typecheck.check_type(arg.internal_representation, anonymous_tuple.AnonymousTuple) py_typecheck.check_len(arg.internal_representation, 2) keys = [k for k, _ in anonymous_tuple.to_elements(arg.type_signature)] vals = [arg.internal_representation[n] for n in [0, 1]] types = [arg.type_signature[n] for n in [0, 1]] for n in [0, 1]: type_analysis.check_federated_type( types[n], placement=placement_literals.CLIENTS) types[n] = type_factory.at_clients(types[n].member) py_typecheck.check_type(vals[n], list) py_typecheck.check_len(vals[n], len(self._target_executors)) item_type = computation_types.StructType([ ((keys[n], types[n].member) if keys[n] else types[n].member) for n in [0, 1] ]) result_type = type_factory.at_clients(item_type) zip_type = computation_types.FunctionType( computation_types.StructType([ ((keys[n], types[n]) if keys[n] else types[n]) for n in [0, 1] ]), result_type) zip_comp = executor_utils.create_intrinsic_comp( intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS, zip_type) async def _child_fn(ex, x, y): py_typecheck.check_type(x, executor_value_base.ExecutorValue) py_typecheck.check_type(y, executor_value_base.ExecutorValue) return await ex.create_call( await ex.create_value(zip_comp, zip_type), await ex.create_struct( anonymous_tuple.AnonymousTuple([(keys[0], x), (keys[1], y)]))) result = await asyncio.gather(*[ _child_fn(c, x, y) for c, x, y in zip(self._target_executors, vals[0], vals[1]) ]) return FederatedComposingStrategyValue(result, result_type)
def create_dummy_intrinsic_def_federated_aggregate(): value = intrinsic_defs.FEDERATED_AGGREGATE type_signature = computation_types.FunctionType([ type_factory.at_clients(tf.float32), tf.float32, type_factory.reduction_op(tf.float32, tf.float32), type_factory.binary_op(tf.float32), computation_types.FunctionType(tf.float32, tf.float32), ], type_factory.at_server(tf.float32)) return value, type_signature
def test_federated_mean(self): @computations.federated_computation(type_factory.at_clients(tf.float32) ) def comp(x): return intrinsics.federated_mean(x) executor, num_clients = _create_test_executor() arg = [float(x + 1) for x in range(num_clients)] result = _invoke(executor, comp, arg) self.assertEqual(result, 6.5)
def test_execution_with_inferred_clients_larger_than_fanout( self, executor_factory_fn): @computations.federated_computation(type_factory.at_clients(tf.int32)) def foo(x): return intrinsics.federated_sum(x) executor = executor_factory_fn(max_fanout=3) with executor_test_utils.install_executor(executor): result = foo([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) self.assertEqual(result, 55)
def test_returns_value_with_federated_type_at_clients(self): value = [ eager_tf_executor.EagerValue(10.0, None, tf.float32), eager_tf_executor.EagerValue(11.0, None, tf.float32), eager_tf_executor.EagerValue(12.0, None, tf.float32), ] type_signature = type_factory.at_clients(tf.float32) value = federating_executor.FederatingExecutorValue(value, type_signature) result = self.run_sync(value.compute()) self.assertEqual(result, [10.0, 11.0, 12.0])
async def _map(self, arg, all_equal=None): py_typecheck.check_type(arg.internal_representation, anonymous_tuple.AnonymousTuple) py_typecheck.check_len(arg.internal_representation, 2) fn_type = arg.type_signature[0] py_typecheck.check_type(fn_type, computation_types.FunctionType) val_type = arg.type_signature[1] py_typecheck.check_type(val_type, computation_types.FederatedType) if all_equal is None: all_equal = val_type.all_equal elif all_equal and not val_type.all_equal: raise ValueError( 'Cannot map a non-all_equal argument into an all_equal result.' ) fn = arg.internal_representation[0] py_typecheck.check_type(fn, pb.Computation) val = arg.internal_representation[1] py_typecheck.check_type(val, list) map_type = computation_types.FunctionType( [fn_type, type_factory.at_clients(fn_type.parameter)], type_factory.at_clients(fn_type.result)) map_comp = executor_utils.create_intrinsic_comp( intrinsic_defs.FEDERATED_MAP, map_type) async def _child_fn(ex, v): py_typecheck.check_type(v, executor_value_base.ExecutorValue) fn_val = await ex.create_value(fn, fn_type) map_val, map_arg = await asyncio.gather( ex.create_value(map_comp, map_type), ex.create_struct([fn_val, v])) return await ex.create_call(map_val, map_arg) result_vals = await asyncio.gather( *[_child_fn(c, v) for c, v in zip(self._target_executors, val)]) federated_type = computation_types.FederatedType(fn_type.result, val_type.placement, all_equal=all_equal) return FederatedComposingStrategyValue(result_vals, federated_type)
def test_returns_value_with_unplaced_type_and_clients(self, executor): value, type_signature = executor_test_utils.create_dummy_value_unplaced() value = self.run_sync(executor.create_value(value, type_signature)) result = self.run_sync( executor_utils.compute_intrinsic_federated_value( executor, value, placement_literals.CLIENTS)) self.assertIsInstance(result, executor_value_base.ExecutorValue) expected_type = type_factory.at_clients(type_signature, all_equal=True) self.assertEqual(result.type_signature.compact_representation(), expected_type.compact_representation()) actual_result = self.run_sync(result.compute()) self.assertEqual(actual_result, 10.0)
def test_recovers_from_raising(self): class _RaisingExecutor(eager_tf_executor.EagerTFExecutor): """An executor which can be configured to raise on `create_value`.""" def __init__(self): self._should_raise = True super().__init__() def stop_raising(self): self._should_raise = False async def create_value(self, *args, **kwargs): if self._should_raise: raise AssertionError return await super().create_value(*args, **kwargs) raising_executors = [_RaisingExecutor() for _ in range(2)] factory = federated_resolving_strategy.FederatedResolvingStrategy.factory( { placement_literals.SERVER: _create_worker_stack(), placement_literals.CLIENTS: raising_executors, }) federating_ex = federating_executor.FederatingExecutor( factory, _create_worker_stack()) raising_stacks = [federating_ex for _ in range(3)] executor = _create_middle_stack([ _create_middle_stack(raising_stacks), _create_middle_stack(raising_stacks), ]) @computations.federated_computation(type_factory.at_clients(tf.float32) ) def comp(x): return intrinsics.federated_mean(x) # 2 clients per worker stack * 3 worker stacks * 2 middle stacks num_clients = 12 arg = [float(x + 1) for x in range(num_clients)] with self.assertRaises(AssertionError): _invoke(executor, comp, arg) for ex in raising_executors: ex.stop_raising() result = _invoke(executor, comp, arg) self.assertEqual(result, 6.5)
def test_changing_cardinalities_across_calls(self): @computations.federated_computation(type_factory.at_clients(tf.int32)) def comp(x): return x five_ints = list(range(5)) ten_ints = list(range(10)) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): five = comp(five_ints) ten = comp(ten_ints) self.assertEqual(five, five_ints) self.assertEqual(ten, ten_ints)
def test_returns_value_with_federated_type_at_server(self, executor, num_clients): del num_clients # Unused. value, type_signature = executor_test_utils.create_dummy_value_at_server() value = self.run_sync(executor.create_value(value, type_signature)) result = self.run_sync( executor_utils.compute_intrinsic_federated_broadcast(executor, value)) self.assertIsInstance(result, executor_value_base.ExecutorValue) expected_type = type_factory.at_clients( type_signature.member, all_equal=True) self.assertEqual(result.type_signature.compact_representation(), expected_type.compact_representation()) actual_result = self.run_sync(result.compute()) self.assertEqual(actual_result, 10.0)
def test_get_size_info(self, num_clients): @computations.federated_computation( type_factory.at_clients(computation_types.SequenceType( tf.float32)), type_factory.at_server(tf.float32)) def comp(temperatures, threshold): client_data = [ temperatures, intrinsics.federated_broadcast(threshold) ] result_map = intrinsics.federated_map( count_over, intrinsics.federated_zip(client_data)) count_map = intrinsics.federated_map(count_total, temperatures) return intrinsics.federated_mean(result_map, count_map) sizing_factory = executor_stacks.sizing_executor_factory( num_clients=num_clients) sizing_context = execution_context.ExecutionContext(sizing_factory) with get_context_stack.get_context_stack().install(sizing_context): to_float = lambda x: tf.cast(x, tf.float32) temperatures = [tf.data.Dataset.range(10).map(to_float) ] * num_clients threshold = 15.0 comp(temperatures, threshold) # Each client receives a tf.float32 and uploads two tf.float32 values. expected_broadcast_bits = [num_clients * 32] expected_aggregate_bits = [num_clients * 32 * 2] expected_broadcast_history = { (('CLIENTS', num_clients), ): [[1, tf.float32]] * num_clients } expected_aggregate_history = { (('CLIENTS', num_clients), ): [[1, tf.float32]] * num_clients * 2 } size_info = sizing_factory.get_size_info() self.assertEqual(expected_broadcast_history, size_info.broadcast_history) self.assertEqual(expected_aggregate_history, size_info.aggregate_history) self.assertEqual(expected_broadcast_bits, size_info.broadcast_bits) self.assertEqual(expected_aggregate_bits, size_info.aggregate_bits)
def _create_tff_parallel_clients_with_dataset_reduce(): @tf.function def reduce_fn(x, y): return x + y @tf.function def dataset_reduce_fn(ds, initial_val): return ds.reduce(initial_val, reduce_fn) @computations.tf_computation(computation_types.SequenceType(tf.int64)) def dataset_reduce_fn_wrapper(ds): initial_val = tf.Variable(np.int64(1.0)) return dataset_reduce_fn(ds, initial_val) @computations.federated_computation( type_factory.at_clients(computation_types.SequenceType(tf.int64))) def parallel_client_run(client_datasets): return intrinsics.federated_map(dataset_reduce_fn_wrapper, client_datasets) return parallel_client_run
async def _compute_intrinsic_federated_mean(self, arg): member_type = arg.type_signature.member ones = await self.create_value( 1, type_factory.at_clients(member_type, all_equal=True)) totals = (await self._compute_intrinsic_federated_sum( await self._compute_intrinsic_federated_zip_at_clients( await self.create_tuple([arg, ones])))).internal_representation py_typecheck.check_type(totals, executor_value_base.ExecutorValue) fed_sum, count = tuple(await asyncio.gather( self._parent_executor.create_selection(totals, index=0), self._parent_executor.create_selection(totals, index=1))) count_val = await count.compute() factor, multiply = tuple(await asyncio.gather(*[ executor_utils.embed_tf_scalar_constant( self._parent_executor, member_type, float(1.0 / count_val)), executor_utils.embed_tf_binary_operator(self._parent_executor, member_type, tf.multiply) ])) multiply_arg = await self._parent_executor.create_tuple( [fed_sum, factor]) result = await self._parent_executor.create_call( multiply, multiply_arg) return CompositeValue(result, type_factory.at_server(member_type))
def create_dummy_value_at_clients_all_equal(): """Returns a Python value and federated type at clients and all equal.""" value = 10.0 type_signature = type_factory.at_clients(tf.float32, all_equal=True) return value, type_signature
def create_dummy_value_at_clients(number_of_clients: int = 3): """Returns a Python value and federated type at clients.""" value = [float(x) for x in range(10, number_of_clients + 10)] type_signature = type_factory.at_clients(tf.float32) return value, type_signature
def create_dummy_intrinsic_def_federated_value_at_clients(): value = intrinsic_defs.FEDERATED_VALUE_AT_CLIENTS type_signature = computation_types.FunctionType( tf.float32, type_factory.at_clients(tf.float32, all_equal=True)) return value, type_signature
def create_dummy_intrinsic_def_federated_sum(): value = intrinsic_defs.FEDERATED_SUM type_signature = computation_types.FunctionType( type_factory.at_clients(tf.float32), type_factory.at_server(tf.float32)) return value, type_signature
def create_dummy_intrinsic_def_federated_eval_at_clients(): value = intrinsic_defs.FEDERATED_EVAL_AT_CLIENTS type_signature = computation_types.FunctionType( computation_types.FunctionType(None, tf.float32), type_factory.at_clients(tf.float32)) return value, type_signature
def create_dummy_intrinsic_def_federated_collect(): value = intrinsic_defs.FEDERATED_COLLECT type_signature = computation_types.FunctionType( type_factory.at_clients(tf.float32), type_factory.at_server(computation_types.SequenceType(tf.float32))) return value, type_signature
def create_dummy_intrinsic_def_federated_broadcast(): value = intrinsic_defs.FEDERATED_BROADCAST type_signature = computation_types.FunctionType( type_factory.at_server(tf.float32), type_factory.at_clients(tf.float32, all_equal=True)) return value, type_signature
def test_at_clients(self): type_spec = computation_types.TensorType(tf.bool) actual_type = type_factory.at_clients(type_spec) expected_type = computation_types.FederatedType( type_spec, placement_literals.CLIENTS) self.assertEqual(actual_type, expected_type)
# intrinsics defined above, as follows. # # @federated_computation # def federated_aggregate(x, zero, accumulate, merge, report): # a = generic_partial_reduce(x, zero, accumulate, INTERMEDIATE_AGGREGATORS) # b = generic_reduce(a, zero, merge, SERVER) # c = generic_map(report, b) # return c # # Actual implementations might vary. # # Type signature: <{T}@CLIENTS,U,(<U,T>->U),(<U,U>->U),(U->R)> -> R@SERVER FEDERATED_AGGREGATE = IntrinsicDef( 'FEDERATED_AGGREGATE', 'federated_aggregate', computation_types.FunctionType(parameter=[ type_factory.at_clients(computation_types.AbstractType('T')), computation_types.AbstractType('U'), type_factory.reduction_op(computation_types.AbstractType('U'), computation_types.AbstractType('T')), type_factory.binary_op(computation_types.AbstractType('U')), computation_types.FunctionType(computation_types.AbstractType('U'), computation_types.AbstractType('R')) ], result=type_factory.at_server( computation_types.AbstractType('R')))) # Applies a given function to a value on the server. # # Type signature: <(T->U),T@SERVER> -> U@SERVER FEDERATED_APPLY = IntrinsicDef( 'FEDERATED_APPLY', 'federated_apply',
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([ type_factory.at_clients(arg.type_signature[0].member), type_factory.at_clients(arg.type_signature[1].member) ]), type_factory.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]), type_factory.at_clients(multiply_blk.type_signature.result)) sum1_type = computation_types.FunctionType( type_factory.at_clients(map_type.result.member), type_factory.at_server(map_type.result.member)) sum2_type = computation_types.FunctionType( type_factory.at_clients(arg.type_signature[1].member), type_factory.at_server(arg.type_signature[1].member)) zip2_type = computation_types.FunctionType( computation_types.StructType([sum1_type.result, sum2_type.result]), type_factory.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]), type_factory.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 ) -> 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.NamedTupleType([ type_factory.at_clients(arg.type_signature[0].member), type_factory.at_clients(arg.type_signature[1].member) ]), type_factory.at_clients( computation_types.NamedTupleType( [arg.type_signature[0].member, arg.type_signature[1].member]))) zip1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS, zip1_type) zipped_arg = await executor.create_call( await executor.create_value(zip1_comp, zip1_type), arg) # TODO(b/134543154): Replace with something that produces a section of # plain TensorFlow code instead of constructing a lambda (so that this # can be executed directly on top of a plain TensorFlow-based executor). multiply_blk = building_block_factory.create_binary_operator_with_upcast( zipped_arg.type_signature.member, tf.multiply) map_type = computation_types.FunctionType( computation_types.NamedTupleType( [multiply_blk.type_signature, zipped_arg.type_signature]), type_factory.at_clients(multiply_blk.type_signature.result)) map_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_MAP, map_type) products = await executor.create_call( await executor.create_value(map_comp, map_type), await executor.create_tuple([ await executor.create_value(multiply_blk.proto, multiply_blk.type_signature), zipped_arg ])) sum1_type = computation_types.FunctionType( type_factory.at_clients(products.type_signature.member), type_factory.at_server(products.type_signature.member)) sum1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM, sum1_type) sum_of_products = await executor.create_call( await executor.create_value(sum1_comp, sum1_type), products) sum2_type = computation_types.FunctionType( type_factory.at_clients(arg.type_signature[1].member), type_factory.at_server(arg.type_signature[1].member)) sum2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM, sum2_type) total_weight = await executor.create_call( *(await asyncio.gather(executor.create_value(sum2_comp, sum2_type), executor.create_selection(arg, index=1)))) zip2_type = computation_types.FunctionType( computation_types.NamedTupleType( [sum_of_products.type_signature, total_weight.type_signature]), type_factory.at_server( computation_types.NamedTupleType([ sum_of_products.type_signature.member, total_weight.type_signature.member ]))) zip2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_ZIP_AT_SERVER, zip2_type) divide_arg = await executor.create_call(*(await asyncio.gather( executor.create_value(zip2_comp, zip2_type), executor.create_tuple([sum_of_products, total_weight])))) divide_blk = building_block_factory.create_binary_operator_with_upcast( divide_arg.type_signature.member, tf.divide) apply_type = computation_types.FunctionType( computation_types.NamedTupleType( [divide_blk.type_signature, divide_arg.type_signature]), type_factory.at_server(divide_blk.type_signature.result)) apply_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_APPLY, apply_type) return await executor.create_call(*(await asyncio.gather( executor.create_value(apply_comp, apply_type), executor.create_tuple([ await executor.create_value(divide_blk.proto, divide_blk.type_signature), divide_arg ]))))
def test_raises_type_error(self): type_analysis.check_valid_federated_weighted_mean_argument_tuple_type( computation_types.StructType([type_factory.at_clients(tf.float32)] * 2)) with self.assertRaises(TypeError): type_analysis.check_valid_federated_weighted_mean_argument_tuple_type( computation_types.StructType([type_factory.at_clients(tf.int32)] * 2))
def test_at_clients(self): self.assertEqual(str(type_factory.at_clients(tf.bool)), '{bool}@CLIENTS')