def test_with_temperature_sensor_example(self): @computations.tf_computation(computation_types.SequenceType( tf.float32), tf.float32) def count_over(ds, t): return ds.reduce( np.float32(0), lambda n, x: n + tf.cast(tf.greater(x, t), tf.float32)) @computations.tf_computation(computation_types.SequenceType(tf.float32) ) def count_total(ds): return ds.reduce(np.float32(0.0), lambda n, _: n + 1.0) @computations.federated_computation( type_factory.at_clients(computation_types.SequenceType( tf.float32)), type_factory.at_server(tf.float32)) def comp(temperatures, threshold): return intrinsics.federated_mean( intrinsics.federated_map( count_over, intrinsics.federated_zip([ temperatures, intrinsics.federated_broadcast(threshold) ])), intrinsics.federated_map(count_total, temperatures)) set_default_executor.set_default_executor( executor_stacks.create_local_executor(3)) to_float = lambda x: tf.cast(x, tf.float32) temperatures = [ tf.data.Dataset.range(10).map(to_float), tf.data.Dataset.range(20).map(to_float), tf.data.Dataset.range(30).map(to_float) ] threshold = 15.0 result = comp(temperatures, threshold) self.assertAlmostEqual(result, 8.333, places=3) set_default_executor.set_default_executor( executor_stacks.create_local_executor()) to_float = lambda x: tf.cast(x, tf.float32) temperatures = [ tf.data.Dataset.range(10).map(to_float), tf.data.Dataset.range(20).map(to_float), tf.data.Dataset.range(30).map(to_float) ] threshold = 15.0 result = comp(temperatures, threshold) self.assertAlmostEqual(result, 8.333, places=3) set_default_executor.set_default_executor()
def test_context(rpc_mode='REQUEST_REPLY'): port = portpicker.pick_unused_port() server_pool = logging_pool.pool(max_workers=1) server = grpc.server(server_pool) server.add_insecure_port('[::]:{}'.format(port)) target_executor = executor_stacks.create_local_executor( num_clients=3)(None) tracer = executor_test_utils.TracingExecutor(target_executor) service = executor_service.ExecutorService(tracer) executor_pb2_grpc.add_ExecutorServicer_to_server(service, server) server.start() channel = grpc.insecure_channel('localhost:{}'.format(port)) remote_exec = remote_executor.RemoteExecutor(channel, rpc_mode) executor = lambda_executor.LambdaExecutor(remote_exec) set_default_executor.set_default_executor(executor) try: yield collections.namedtuple('_', 'executor tracer')(executor, tracer) finally: remote_exec.__del__() set_default_executor.set_default_executor() try: channel.close() except AttributeError: pass # Public gRPC channel doesn't support close() finally: server.stop(None)
def test_with_num_clients_larger_than_fanout(self): set_default_executor.set_default_executor( executor_stacks.create_local_executor(max_fanout=3)) @computations.federated_computation(type_factory.at_clients(tf.int32)) def foo(x): return intrinsics.federated_sum(x) self.assertEqual(foo([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]), 55)
def decorator(fn, named_executors=None): if not named_executors: named_executors = [ ('reference', None), ('local', executor_stacks.create_local_executor()), ] named_parameters_decorator = parameterized.named_parameters( named_executors) fn = executor_decorator(fn) fn = named_parameters_decorator(fn) return fn
def test_with_no_args(self): set_default_executor.set_default_executor( executor_stacks.create_local_executor()) @computations.tf_computation def foo(): return tf.constant(10) self.assertEqual(foo(), 10) set_default_executor.set_default_executor()
def test_with_incomplete_temperature_sensor_example(self): @computations.federated_computation( type_constructors.at_clients( computation_types.SequenceType(tf.float32)), type_constructors.at_server(tf.float32)) def comp(temperatures, threshold): @computations.tf_computation( computation_types.SequenceType(tf.float32), tf.float32) def count(ds, t): return ds.reduce( np.int32(0), lambda n, x: n + tf.cast(tf.greater(x, t), tf.int32)) return intrinsics.federated_map( count, intrinsics.federated_zip( [temperatures, intrinsics.federated_broadcast(threshold)])) num_clients = 10 set_default_executor.set_default_executor( executor_stacks.create_local_executor(num_clients)) temperatures = [ tf.data.Dataset.range(1000).map(lambda x: tf.cast(x, tf.float32)) for _ in range(num_clients) ] threshold = 100.0 result = comp(temperatures, threshold) self.assertCountEqual([x.numpy() for x in result], [899 for _ in range(num_clients)]) set_default_executor.set_default_executor()
def test_with_mnist_training_example_unspecified_clients(self): executor_test_utils.test_mnist_training( self, executor_stacks.create_local_executor())
def test_with_mnist_training_example(self): executor_test_utils.test_mnist_training( self, executor_stacks.create_local_executor(1))
def test_raises_with_max_fanout_1(self): with self.assertRaises(ValueError): executor_stacks.create_local_executor(2, 1)
class TensorFlowComputationsWithDatasetsTest(parameterized.TestCase): # TODO(b/122081673): Support tf.Dataset serialization in tf2_computation. def test_with_tf_datasets(self): @tff.tf_computation(tff.SequenceType(tf.int64)) def consume(ds): return ds.reduce(np.int64(0), lambda x, y: x + y) self.assertEqual(str(consume.type_signature), '(int64* -> int64)') @tff.tf_computation def produce(): return tf.data.Dataset.range(10) self.assertEqual(str(produce.type_signature), '( -> int64*)') self.assertEqual(consume(produce()), 45) # TODO(b/131363314): The reference executor should support generating and # returning infinite datasets @core_test.executors( ('local', executor_stacks.create_local_executor(1)),) def test_consume_infinite_tf_dataset(self): @tff.tf_computation(tff.SequenceType(tf.int64)) def consume(ds): # Consume the first 10 elements of the dataset. return ds.take(10).reduce(np.int64(0), lambda x, y: x + y) self.assertEqual(consume(tf.data.Dataset.range(10).repeat()), 45) # TODO(b/131363314): The reference executor should support generating and # returning infinite datasets @core_test.executors( ('local', executor_stacks.create_local_executor(1)),) def test_produce_and_consume_infinite_tf_dataset(self): @tff.tf_computation(tff.SequenceType(tf.int64)) def consume(ds): # Consume the first 10 elements of the dataset. return ds.take(10).reduce(np.int64(0), lambda x, y: x + y) @tff.tf_computation def produce(): # Produce an infinite dataset. return tf.data.Dataset.range(10).repeat() self.assertEqual(consume(produce()), 45) def test_with_sequence_of_pairs(self): pairs = tf.data.Dataset.from_tensor_slices( (list(range(5)), list(range(5, 10)))) @tff.tf_computation def process_pairs(ds): return ds.reduce(0, lambda state, pair: state + pair[0] + pair[1]) self.assertEqual(process_pairs(pairs), 45) def test_tf_comp_with_sequence_inputs_and_outputs_does_not_fail(self): @tff.tf_computation(tff.SequenceType(tf.int32)) def _(x): return x def test_with_four_element_dataset_pipeline(self): @tff.tf_computation def comp1(): return tf.data.Dataset.range(5) @tff.tf_computation(tff.SequenceType(tf.int64)) def comp2(ds): return ds.map(lambda x: tf.cast(x + 1, tf.float32)) @tff.tf_computation(tff.SequenceType(tf.float32)) def comp3(ds): return ds.repeat(5) @tff.tf_computation(tff.SequenceType(tf.float32)) def comp4(ds): return ds.reduce(0.0, lambda x, y: x + y) @tff.tf_computation def comp5(): return comp4(comp3(comp2(comp1()))) self.assertEqual(comp5(), 75.0)
def _make_default_context(): return execution_context.ExecutionContext( executor_stacks.create_local_executor())
_ = next(iter(x[1])) def test_fetch_value_with_empty_structured_dataset_and_tensors(self): def return_dataset(): ds1 = tf.data.Dataset.from_tensor_slices( collections.OrderedDict([('a', [1, 1]), ('b', [1, 1])])) return [tf.constant([0., 0.]), ds1.batch(5).take(0)] executable_return_dataset = computation_impl.ComputationImpl( tensorflow_serialization.serialize_py_fn_as_tf_computation( return_dataset, None, context_stack_impl.context_stack)[0], context_stack_impl.context_stack) x = executable_return_dataset() self.assertAllEqual(x[0], [0., 0.]) self.assertEqual( tf.data.experimental.get_structure(x[1]), collections.OrderedDict([ ('a', tf.TensorSpec(shape=(None, ), dtype=tf.int32)), ('b', tf.TensorSpec(shape=(None, ), dtype=tf.int32)), ])) with self.assertRaises(StopIteration): _ = next(iter(x[1])) if __name__ == '__main__': # Use the local executor. set_default_executor.set_default_executor( executor_stacks.create_local_executor()) test.main()
class IntrinsicsTest(parameterized.TestCase): def test_federated_broadcast_with_server_all_equal_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.SERVER)) def foo(x): return intrinsics.federated_broadcast(x) self.assertEqual(str(foo.type_signature), '(int32@SERVER -> int32@CLIENTS)') def test_federated_broadcast_with_server_non_all_equal_int(self): with self.assertRaises(TypeError): @computations.federated_computation( computation_types.FederatedType( tf.int32, placements.SERVER, all_equal=False)) def _(x): return intrinsics.federated_broadcast(x) def test_federated_broadcast_with_client_int(self): with self.assertRaises(TypeError): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS, True)) def _(x): return intrinsics.federated_broadcast(x) def test_federated_broadcast_with_non_federated_val(self): with self.assertRaises(TypeError): @computations.federated_computation(tf.int32) def _(x): return intrinsics.federated_broadcast(x) def test_federated_map_with_client_all_equal_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS, True)) def foo(x): return intrinsics.federated_map( computations.tf_computation(lambda x: x > 10, tf.int32), x) self.assertEqual( str(foo.type_signature), '(int32@CLIENTS -> {bool}@CLIENTS)') def test_federated_map_with_client_non_all_equal_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): return intrinsics.federated_map( computations.tf_computation(lambda x: x > 10, tf.int32), x) self.assertEqual( str(foo.type_signature), '({int32}@CLIENTS -> {bool}@CLIENTS)') def test_federated_map_with_server_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.SERVER)) def foo(x): return intrinsics.federated_map( computations.tf_computation(lambda x: x > 10, tf.int32), x) self.assertEqual(str(foo.type_signature), '(int32@SERVER -> bool@SERVER)') def test_federated_map_injected_zip_with_server_int(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.SERVER), computation_types.FederatedType(tf.int32, placements.SERVER) ]) def foo(x, y): return intrinsics.federated_map( computations.tf_computation(lambda x, y: x > 10, [tf.int32, tf.int32]), [x, y]) self.assertEqual( str(foo.type_signature), '(<int32@SERVER,int32@SERVER> -> bool@SERVER)') def test_federated_map_injected_zip_fails_different_placements(self): def foo(x, y): return intrinsics.federated_map( computations.tf_computation(lambda x, y: x > 10, [tf.int32, tf.int32]), [x, y]) with self.assertRaisesRegex( TypeError, 'The value to be mapped must be a FederatedType or implicitly ' 'convertible to a FederatedType.'): computations.federated_computation(foo, [ computation_types.FederatedType(tf.int32, placements.SERVER), computation_types.FederatedType(tf.int32, placements.CLIENTS) ]) def test_federated_map_with_non_federated_val(self): with self.assertRaises(TypeError): @computations.federated_computation(tf.int32) def _(x): return intrinsics.federated_map( computations.tf_computation(lambda x: x > 10, tf.int32), x) def test_federated_sum_with_client_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): return intrinsics.federated_sum(x) self.assertEqual( str(foo.type_signature), '({int32}@CLIENTS -> int32@SERVER)') def test_federated_sum_with_client_string(self): with self.assertRaises(TypeError): @computations.federated_computation( computation_types.FederatedType(tf.string, placements.CLIENTS)) def _(x): return intrinsics.federated_sum(x) def test_federated_sum_with_server_int(self): with self.assertRaises(TypeError): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.SERVER)) def _(x): return intrinsics.federated_sum(x) def test_federated_zip_with_client_non_all_equal_int_and_bool(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.CLIENTS), computation_types.FederatedType(tf.bool, placements.CLIENTS, True) ]) def foo(x, y): return intrinsics.federated_zip([x, y]) self.assertEqual( str(foo.type_signature), '(<{int32}@CLIENTS,bool@CLIENTS> -> {<int32,bool>}@CLIENTS)') def test_federated_zip_with_single_unnamed_int_client(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.CLIENTS), ]) def foo(x): return intrinsics.federated_zip(x) self.assertEqual( str(foo.type_signature), '(<{int32}@CLIENTS> -> {<int32>}@CLIENTS)') def test_federated_zip_with_single_unnamed_int_server(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.SERVER), ]) def foo(x): return intrinsics.federated_zip(x) self.assertEqual( str(foo.type_signature), '(<int32@SERVER> -> <int32>@SERVER)') def test_federated_zip_with_single_named_bool_clients(self): @computations.federated_computation([ ('a', computation_types.FederatedType(tf.bool, placements.CLIENTS)), ]) def foo(x): return intrinsics.federated_zip(x) self.assertEqual( str(foo.type_signature), '(<a={bool}@CLIENTS> -> {<a=bool>}@CLIENTS)') def test_federated_zip_with_single_named_bool_server(self): @computations.federated_computation([ ('a', computation_types.FederatedType(tf.bool, placements.SERVER)), ]) def foo(x): return intrinsics.federated_zip(x) self.assertEqual( str(foo.type_signature), '(<a=bool@SERVER> -> <a=bool>@SERVER)') def test_federated_zip_with_names_client_non_all_equal_int_and_bool(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.CLIENTS), computation_types.FederatedType(tf.bool, placements.CLIENTS, True) ]) def foo(x, y): a = {'x': x, 'y': y} return intrinsics.federated_zip(a) self.assertEqual( str(foo.type_signature), '(<{int32}@CLIENTS,bool@CLIENTS> -> {<x=int32,y=bool>}@CLIENTS)') def test_federated_zip_with_client_all_equal_int_and_bool(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.CLIENTS, True), computation_types.FederatedType(tf.bool, placements.CLIENTS, True) ]) def foo(x, y): return intrinsics.federated_zip([x, y]) self.assertEqual( str(foo.type_signature), '(<int32@CLIENTS,bool@CLIENTS> -> {<int32,bool>}@CLIENTS)') def test_federated_zip_with_names_client_all_equal_int_and_bool(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.CLIENTS, True), computation_types.FederatedType(tf.bool, placements.CLIENTS, True) ]) def foo(arg): a = {'x': arg[0], 'y': arg[1]} return intrinsics.federated_zip(a) self.assertEqual( str(foo.type_signature), '(<int32@CLIENTS,bool@CLIENTS> -> {<x=int32,y=bool>}@CLIENTS)') def test_federated_zip_with_server_int_and_bool(self): @computations.federated_computation([ computation_types.FederatedType(tf.int32, placements.SERVER), computation_types.FederatedType(tf.bool, placements.SERVER) ]) def foo(x, y): return intrinsics.federated_zip([x, y]) self.assertEqual( str(foo.type_signature), '(<int32@SERVER,bool@SERVER> -> <int32,bool>@SERVER)') def test_federated_zip_with_names_server_int_and_bool(self): @computations.federated_computation([ ('a', computation_types.FederatedType(tf.int32, placements.SERVER)), ('b', computation_types.FederatedType(tf.bool, placements.SERVER)), ]) def foo(arg): return intrinsics.federated_zip(arg) self.assertEqual( str(foo.type_signature), '(<a=int32@SERVER,b=bool@SERVER> -> <a=int32,b=bool>@SERVER)') def test_federated_zip_error_different_placements(self): with self.assertRaises(TypeError): @computations.federated_computation([ ('a', computation_types.FederatedType(tf.int32, placements.SERVER)), ('b', computation_types.FederatedType(tf.bool, placements.CLIENTS)), ]) def _(arg): return intrinsics.federated_zip(arg) def test_federated_collect_with_client_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): return intrinsics.federated_collect(x) self.assertEqual( str(foo.type_signature), '({int32}@CLIENTS -> int32*@SERVER)') def test_federated_collect_with_server_int_fails(self): with self.assertRaises(TypeError): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.SERVER)) def _(x): return intrinsics.federated_collect(x) def test_federated_mean_with_client_float32_without_weight(self): @computations.federated_computation( computation_types.FederatedType(tf.float32, placements.CLIENTS)) def foo(x): return intrinsics.federated_mean(x) self.assertEqual( str(foo.type_signature), '({float32}@CLIENTS -> float32@SERVER)') def test_federated_mean_with_all_equal_client_float32_without_weight(self): federated_all_equal_float = computation_types.FederatedType( tf.float32, placements.CLIENTS, all_equal=True) @computations.federated_computation(federated_all_equal_float) def foo(x): return intrinsics.federated_mean(x) self.assertEqual( str(foo.type_signature), '(float32@CLIENTS -> float32@SERVER)') def test_federated_mean_with_all_equal_client_float32_with_weight(self): federated_all_equal_float = computation_types.FederatedType( tf.float32, placements.CLIENTS, all_equal=True) @computations.federated_computation(federated_all_equal_float) def foo(x): return intrinsics.federated_mean(x, x) self.assertEqual( str(foo.type_signature), '(float32@CLIENTS -> float32@SERVER)') def test_federated_mean_with_client_tuple_with_int32_weight(self): @computations.federated_computation([ computation_types.FederatedType([('x', tf.float64), ('y', tf.float64)], placements.CLIENTS), computation_types.FederatedType(tf.int32, placements.CLIENTS) ]) def foo(x, y): return intrinsics.federated_mean(x, y) self.assertEqual( str(foo.type_signature), '(<{<x=float64,y=float64>}@CLIENTS,{int32}@CLIENTS> ' '-> <x=float64,y=float64>@SERVER)') def test_federated_mean_with_client_int32_fails(self): with self.assertRaises(TypeError): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def _(x): return intrinsics.federated_mean(x) def test_federated_mean_with_string_weight_fails(self): with self.assertRaises(TypeError): @computations.federated_computation([ computation_types.FederatedType(tf.float32, placements.CLIENTS), computation_types.FederatedType(tf.string, placements.CLIENTS) ]) def _(x, y): return intrinsics.federated_mean(x, y) def test_federated_aggregate_with_client_int(self): # The representation used during the aggregation process will be a named # tuple with 2 elements - the integer 'total' that represents the sum of # elements encountered, and the integer element 'count'. # pylint: disable=invalid-name Accumulator = collections.namedtuple('Accumulator', 'total count') # pylint: enable=invalid-name accumulator_type = computation_types.NamedTupleType( Accumulator(tf.int32, tf.int32)) # The operator to use during the first stage simply adds an element to the # total and updates the count. @computations.tf_computation([accumulator_type, tf.int32]) def accumulate(accu, elem): return Accumulator(accu.total + elem, accu.count + 1) # The operator to use during the second stage simply adds total and count. @computations.tf_computation([accumulator_type, accumulator_type]) def merge(x, y): return Accumulator(x.total + y.total, x.count + y.count) # The operator to use during the final stage simply computes the ratio. @computations.tf_computation(accumulator_type) def report(accu): return tf.cast(accu.total, tf.float32) / tf.cast(accu.count, tf.float32) @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): return intrinsics.federated_aggregate(x, Accumulator(0, 0), accumulate, merge, report) self.assertEqual( str(foo.type_signature), '({int32}@CLIENTS -> float32@SERVER)') def test_federated_aggregate_with_federated_zero_fails(self): @computations.federated_computation() def build_federated_zero(): return intrinsics.federated_value(0, placements.SERVER) @computations.tf_computation([tf.int32, tf.int32]) def accumulate(accu, elem): return accu + elem # The operator to use during the second stage simply adds total and count. @computations.tf_computation([tf.int32, tf.int32]) def merge(x, y): return x + y # The operator to use during the final stage simply computes the ratio. @computations.tf_computation(tf.int32) def report(accu): return accu def foo(x): return intrinsics.federated_aggregate(x, build_federated_zero(), accumulate, merge, report) with self.assertRaisesRegex( TypeError, 'Expected `zero` to be assignable to type int32, ' 'but was of incompatible type int32@SERVER'): computations.federated_computation( foo, computation_types.FederatedType(tf.int32, placements.CLIENTS)) def test_federated_aggregate_with_unknown_dimension(self): Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name accumulator_type = computation_types.NamedTupleType( Accumulator( samples=computation_types.TensorType(dtype=tf.int32, shape=[None]))) @computations.tf_computation() def build_empty_accumulator(): return Accumulator(samples=tf.zeros(shape=[0], dtype=tf.int32)) # The operator to use during the first stage simply adds an element to the # tensor, increasing its size. @computations.tf_computation([accumulator_type, tf.int32]) def accumulate(accu, elem): return Accumulator( samples=tf.concat( [accu.samples, tf.expand_dims(elem, axis=0)], axis=0)) # The operator to use during the second stage simply adds total and count. @computations.tf_computation([accumulator_type, accumulator_type]) def merge(x, y): return Accumulator(samples=tf.concat([x.samples, y.samples], axis=0)) # The operator to use during the final stage simply computes the ratio. @computations.tf_computation(accumulator_type) def report(accu): return accu @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): return intrinsics.federated_aggregate(x, build_empty_accumulator(), accumulate, merge, report) self.assertEqual( str(foo.type_signature), '({int32}@CLIENTS -> <samples=int32[?]>@SERVER)') def test_federated_reduce_with_tf_add_raw_constant(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): plus = computations.tf_computation(tf.add, [tf.int32, tf.int32]) return intrinsics.federated_reduce(x, 0, plus) self.assertEqual( str(foo.type_signature), '({int32}@CLIENTS -> int32@SERVER)') def test_num_over_temperature_threshold_example(self): @computations.federated_computation([ computation_types.FederatedType(tf.float32, placements.CLIENTS), computation_types.FederatedType(tf.float32, placements.SERVER) ]) def foo(temperatures, threshold): return intrinsics.federated_sum( intrinsics.federated_map( computations.tf_computation( lambda x, y: tf.cast(tf.greater(x, y), tf.int32), [tf.float32, tf.float32]), [temperatures, intrinsics.federated_broadcast(threshold)])) self.assertEqual( str(foo.type_signature), '(<{float32}@CLIENTS,float32@SERVER> -> int32@SERVER)') @parameterized.named_parameters(('test_n_2', 2), ('test_n_3', 3), ('test_n_5', 5)) def test_n_tuple_federated_zip_tensor_args(self, n): fed_type = computation_types.FederatedType(tf.int32, placements.CLIENTS) initial_tuple_type = computation_types.NamedTupleType([fed_type] * n) final_fed_type = computation_types.FederatedType([tf.int32] * n, placements.CLIENTS) function_type = computation_types.FunctionType(initial_tuple_type, final_fed_type) type_string = str(function_type) @computations.federated_computation( [computation_types.FederatedType(tf.int32, placements.CLIENTS)] * n) def foo(x): return intrinsics.federated_zip(x) self.assertEqual(str(foo.type_signature), type_string) @parameterized.named_parameters( ('test_n_2_int', 2, computation_types.FederatedType(tf.int32, placements.CLIENTS)), ('test_n_3_int', 3, computation_types.FederatedType(tf.int32, placements.CLIENTS)), ('test_n_5_int', 5, computation_types.FederatedType(tf.int32, placements.CLIENTS)), ('test_n_2_tuple', 2, computation_types.FederatedType([tf.int32, tf.int32], placements.CLIENTS)), ('test_n_3_tuple', 3, computation_types.FederatedType([tf.int32, tf.int32], placements.CLIENTS)), ('test_n_5_tuple', 5, computation_types.FederatedType([tf.int32, tf.int32], placements.CLIENTS))) def test_named_n_tuple_federated_zip(self, n, fed_type): initial_tuple_type = computation_types.NamedTupleType([fed_type] * n) named_fed_type = computation_types.FederatedType( [(str(k), fed_type.member) for k in range(n)], placements.CLIENTS) mixed_fed_type = computation_types.FederatedType( [(str(k), fed_type.member) if k % 2 == 0 else fed_type.member for k in range(n)], placements.CLIENTS) named_function_type = computation_types.FunctionType( initial_tuple_type, named_fed_type) mixed_function_type = computation_types.FunctionType( initial_tuple_type, mixed_fed_type) named_type_string = str(named_function_type) mixed_type_string = str(mixed_function_type) @computations.federated_computation([fed_type] * n) def foo(x): arg = {str(k): x[k] for k in range(n)} return intrinsics.federated_zip(arg) self.assertEqual(str(foo.type_signature), named_type_string) def _make_test_tuple(x, k): """Make a test tuple with a name if k is even, otherwise unnamed.""" if k % 2 == 0: return str(k), x[k] else: return None, x[k] @computations.federated_computation([fed_type] * n) def bar(x): arg = anonymous_tuple.AnonymousTuple( _make_test_tuple(x, k) for k in range(n)) return intrinsics.federated_zip(arg) self.assertEqual(str(bar.type_signature), mixed_type_string) @parameterized.named_parameters([ ('test_n_' + str(n) + '_m_' + str(m), n, m) for n, m in itertools.product([1, 2, 3], [1, 2, 3]) ]) def test_n_tuple_federated_zip_mixed_args(self, n, m): tuple_fed_type = computation_types.FederatedType([tf.int32, tf.int32], placements.CLIENTS) single_fed_type = computation_types.FederatedType(tf.int32, placements.CLIENTS) initial_tuple_type = computation_types.NamedTupleType([tuple_fed_type] * n + [single_fed_type] * m) final_fed_type = computation_types.FederatedType( [[tf.int32, tf.int32]] * n + [tf.int32] * m, placements.CLIENTS) function_type = computation_types.FunctionType(initial_tuple_type, final_fed_type) type_string = str(function_type) @computations.federated_computation([ computation_types.FederatedType( computation_types.NamedTupleType([tf.int32, tf.int32]), placements.CLIENTS) ] * n + [computation_types.FederatedType(tf.int32, placements.CLIENTS)] * m) def baz(x): return intrinsics.federated_zip(x) self.assertEqual(str(baz.type_signature), type_string) def test_federated_apply_raises_warning(self): with warnings.catch_warnings(record=True) as w: warnings.simplefilter('always') @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.SERVER)) def foo(x): return intrinsics.federated_apply( computations.tf_computation(lambda x: x * x, tf.int32), x) self.assertLen(w, 1) self.assertIsInstance(w[0].category(), DeprecationWarning) self.assertIn('tff.federated_apply() is deprecated', str(w[0].message)) self.assertEqual( str(foo.type_signature), '(int32@SERVER -> int32@SERVER)') def test_federated_value_with_bool_on_clients(self): @computations.federated_computation(tf.bool) def foo(x): return intrinsics.federated_value(x, placements.CLIENTS) self.assertEqual(str(foo.type_signature), '(bool -> bool@CLIENTS)') def test_federated_value_raw_np_scalar(self): @computations.federated_computation def test_np_values(): floatv = np.float64(0) tff_float = intrinsics.federated_value(floatv, placements.SERVER) self.assertEqual(str(tff_float.type_signature), 'float64@SERVER') intv = np.int64(0) tff_int = intrinsics.federated_value(intv, placements.SERVER) self.assertEqual(str(tff_int.type_signature), 'int64@SERVER') return (tff_float, tff_int) floatv, intv = test_np_values() self.assertEqual(floatv, 0.0) self.assertEqual(intv, 0) def test_federated_value_raw_tf_scalar_variable(self): v = tf.Variable(initial_value=0., name='test_var') with self.assertRaisesRegex( TypeError, 'TensorFlow construct (.*) has been ' 'encountered in a federated context.'): _ = intrinsics.federated_value(v, placements.SERVER) def test_federated_value_with_bool_on_server(self): @computations.federated_computation(tf.bool) def foo(x): return intrinsics.federated_value(x, placements.SERVER) self.assertEqual(str(foo.type_signature), '(bool -> bool@SERVER)') def test_sequence_sum(self): @computations.federated_computation( computation_types.SequenceType(tf.int32)) def foo1(x): return intrinsics.sequence_sum(x) self.assertEqual(str(foo1.type_signature), '(int32* -> int32)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.SERVER)) def foo2(x): return intrinsics.sequence_sum(x) self.assertEqual( str(foo2.type_signature), '(int32*@SERVER -> int32@SERVER)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS)) def foo3(x): return intrinsics.sequence_sum(x) self.assertEqual( str(foo3.type_signature), '({int32*}@CLIENTS -> {int32}@CLIENTS)') def test_sequence_map(self): @computations.tf_computation(tf.int32) def over_threshold(x): return x > 10 @computations.federated_computation( computation_types.SequenceType(tf.int32)) def foo1(x): return intrinsics.sequence_map(over_threshold, x) self.assertEqual(str(foo1.type_signature), '(int32* -> bool*)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.SERVER)) def foo2(x): return intrinsics.sequence_map(over_threshold, x) self.assertEqual( str(foo2.type_signature), '(int32*@SERVER -> bool*@SERVER)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS)) def foo3(x): return intrinsics.sequence_map(over_threshold, x) self.assertEqual( str(foo3.type_signature), '({int32*}@CLIENTS -> {bool*}@CLIENTS)') def test_sequence_reduce(self): add_numbers = computations.tf_computation(tf.add, [tf.int32, tf.int32]) @computations.federated_computation( computation_types.SequenceType(tf.int32)) def foo1(x): return intrinsics.sequence_reduce(x, 0, add_numbers) self.assertEqual(str(foo1.type_signature), '(int32* -> int32)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.SERVER)) def foo2(x): return intrinsics.sequence_reduce(x, 0, add_numbers) self.assertEqual( str(foo2.type_signature), '(int32*@SERVER -> int32@SERVER)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS)) def foo3(x): return intrinsics.sequence_reduce(x, 0, add_numbers) self.assertEqual( str(foo3.type_signature), '({int32*}@CLIENTS -> {int32}@CLIENTS)') @core_test.executors( ('local', executor_stacks.create_local_executor()),) def test_federated_zip_with_twenty_elements_local_executor(self): n = 20 n_clients = 2 @computations.federated_computation( [computation_types.FederatedType(tf.int32, placements.CLIENTS)] * n) def foo(x): return intrinsics.federated_zip(x) data = [list(range(n_clients)) for _ in range(n)] # This would not have ever returned when local executor was scaling # factorially with number of elements zipped foo(data)
def test_runs_tf_unspecified_clients(self): executor_test_utils.test_runs_tf( self, executor_stacks.create_local_executor())
def test_runs_tf(self): executor_test_utils.test_runs_tf( self, executor_stacks.create_local_executor(1))
state, mean = federated_aggregate_test([1.0, 2.0, 3.0]) self.assertAlmostEqual(mean, 2.0) # (1 + 2 + 3) / (1 + 1 + 1) self.assertDictEqual(state._asdict(), {'call_count': 1}) def test_execute_with_explicit_weights(self): aggregate_fn = computation_utils.StatefulAggregateFn( initialize_fn=agg_initialize_fn, next_fn=agg_next_fn) @computations.federated_computation( computation_types.FederatedType(tf.float32, placements.CLIENTS), computation_types.FederatedType(tf.float32, placements.CLIENTS)) def federated_aggregate_test(args, weights): state = intrinsics.federated_value(aggregate_fn.initialize(), placements.SERVER) return aggregate_fn(state, args, weights) state, mean = federated_aggregate_test([1.0, 2.0, 3.0], [4.0, 1.0, 1.0]) self.assertAlmostEqual(mean, 1.5) # (1*4 + 2*1 + 3*1) / (4 + 1 + 1) self.assertDictEqual(state._asdict(), {'call_count': 1}) if __name__ == '__main__': tf.compat.v1.enable_v2_behavior() # NOTE: num_clients must be explicit here to correctly test the broadcast # behavior. Otherwise TFF will infer there are zero clients, which is an # error. set_default_executor.set_default_executor( executor_stacks.create_local_executor(num_clients=3)) absltest.main()
def _test_ctx(num_clients=None): return execution_context.ExecutionContext( executor_stacks.create_local_executor(num_clients))