class ExecutorsTest(parameterized.TestCase): @executor_test_utils.executors def test_without_arguments(self): @computations.tf_computation(tf.int32) def add_one(x): return x + 1 result = add_one(5) self.assertEqual(result, 6) @executor_test_utils.executors() def test_with_no_arguments(self): @computations.tf_computation(tf.int32) def add_one(x): return x + 1 result = add_one(5) self.assertEqual(result, 6) @executor_test_utils.executors( ('local', executor_stacks.local_executor_factory()),) def test_with_one_argument(self): @computations.tf_computation(tf.int32) def add_one(x): return x + 1 result = add_one(5) self.assertEqual(result, 6) @executor_test_utils.executors( ('local', executor_stacks.local_executor_factory()), ('sizing', executor_stacks.sizing_executor_factory()), ) def test_with_two_argument(self): @computations.tf_computation(tf.int32) def add_one(x): return x + 1 result = add_one(5) self.assertEqual(result, 6)
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.local_executor_factory( num_clients=3).create_executor({}) 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 = reference_resolving_executor.ReferenceResolvingExecutor( remote_exec) try: yield collections.namedtuple('_', 'executor tracer')(executor, tracer) finally: executor.close() tracer.close() try: channel.close() except AttributeError: pass # Public gRPC channel doesn't support close() finally: server.stop(None)
def test_in_executor_stack(self): type_spec = computation_types.SequenceType(tf.int64) ex = data_executor.DataExecutor( eager_tf_executor.EagerTFExecutor(), TestDataBackend(self, 'foo://bar', tf.data.Dataset.range(5), type_spec)) ex_fn = lambda device: ex factory = executor_stacks.local_executor_factory( leaf_executor_fn=ex_fn) context = execution_context.ExecutionContext(executor_fn=factory) @computations.tf_computation(type_spec) def foo(ds): return tf.cast(ds.reduce(np.int64(0), lambda p, q: p + q), tf.int32) @computations.federated_computation def bar(): ds = tff_data.data('foo://bar', type_spec) return foo(ds) with context_stack_impl.context_stack.install(context): result = bar() self.assertEqual(result, 10)
def test_one_arg_tf_computation_with_int_param_and_result(self): @computations.tf_computation(tf.int32) def comp(x): return tf.add(x, 10) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp(3) self.assertEqual(result, 13)
def test_simple_no_arg_tf_computation_with_int_result(self): @computations.tf_computation def comp(): return tf.constant(10) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp() self.assertEqual(result, 10)
def test_three_arg_tf_computation_with_int_params_and_result(self): @computations.tf_computation(tf.int32, tf.int32, tf.int32) def comp(x, y, z): return tf.multiply(tf.add(x, y), z) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp(3, 4, 5) self.assertEqual(result, 35)
def test_tf_computation_with_dataset_params_and_int_result(self): @computations.tf_computation(computation_types.SequenceType(tf.int32)) def comp(ds): return ds.reduce(np.int32(0), lambda x, y: x + y) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): ds = tf.data.Dataset.range(10).map(lambda x: tf.cast(x, tf.int32)) result = comp(ds) self.assertEqual(result, 45)
def _create_concurrent_maxthread_tuples(): tuples = [] for concurrency in range(1, 5): local_ex_string = 'local_executor_{}_client_thread'.format(concurrency) ex_factory = executor_stacks.local_executor_factory( num_client_executors=concurrency) tuples.append((local_ex_string, ex_factory, concurrency)) sizing_ex_string = 'sizing_executor_{}_client_thread'.format(concurrency) ex_factory = executor_stacks.sizing_executor_factory( num_client_executors=concurrency) tuples.append((sizing_ex_string, ex_factory, concurrency)) return tuples
def test_tuple_argument_can_accept_unnamed_elements(self): @computations.tf_computation(tf.int32, tf.int32) def foo(x, y): return x + y executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): # pylint:disable=no-value-for-parameter result = foo(structure.Struct([(None, 2), (None, 3)])) # pylint:enable=no-value-for-parameter self.assertEqual(result, 5)
def decorator(fn, *named_executors): """Construct a custom `parameterized.named_parameter` decorator for `fn`.""" if not named_executors: named_executors = [ ('reference', reference_executor.ReferenceExecutor(compiler=None)), ('local', executor_stacks.local_executor_factory()), ] named_parameters_decorator = parameterized.named_parameters( *named_executors) fn = executor_decorator(fn) fn = named_parameters_decorator(fn) return fn
def test_tf_computation_with_structured_result(self): @computations.tf_computation def comp(): return collections.OrderedDict([ ('a', tf.constant(10)), ('b', tf.constant(20)), ]) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp() self.assertIsInstance(result, collections.OrderedDict) self.assertDictEqual(result, {'a': 10, 'b': 20})
def create_local_execution_context(num_clients=None, max_fanout=100, clients_per_thread=1, server_tf_device=None, client_tf_devices=tuple()): """Creates an execution context that executes computations locally.""" factory = executor_stacks.local_executor_factory( num_clients=num_clients, max_fanout=max_fanout, clients_per_thread=clients_per_thread, server_tf_device=server_tf_device, client_tf_devices=client_tf_devices) return execution_context.ExecutionContext( executor_fn=factory, compiler_fn=compiler.transform_to_native_form)
def test_local_executor_multi_gpus_iter_dataset(self, tf_device): tf_devices = tf.config.list_logical_devices(tf_device) server_tf_device = None if not tf_devices else tf_devices[0] gpu_devices = tf.config.list_logical_devices('GPU') local_executor = executor_stacks.local_executor_factory( server_tf_device=server_tf_device, client_tf_devices=gpu_devices) with executor_test_utils.install_executor(local_executor): parallel_client_run = _create_tff_parallel_clients_with_iter_dataset() client_data = [ tf.data.Dataset.range(10), tf.data.Dataset.range(10).map(lambda x: x + 1) ] client_results = parallel_client_run(client_data) self.assertEqual(client_results, [np.int64(46), np.int64(56)])
def decorator(fn, *named_executors): if not named_executors: named_executors = [ ('local', executor_stacks.local_executor_factory()), ('sizing', executor_stacks.sizing_executor_factory()), ] @parameterized.named_parameters(*named_executors) def wrapped_fn(self, executor): """Install a particular execution context before running `fn`.""" context = execution_context.ExecutionContext(executor) with context_stack_impl.context_stack.install(context): fn(self) return wrapped_fn
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 set_local_execution_context(num_clients=None, max_fanout=100, num_client_executors=32, server_tf_device=None, client_tf_devices=tuple()): """Sets an execution context that executes computations locally.""" factory = executor_stacks.local_executor_factory( num_clients=num_clients, max_fanout=max_fanout, num_client_executors=num_client_executors, server_tf_device=server_tf_device, client_tf_devices=client_tf_devices) context = execution_context.ExecutionContext( executor_fn=factory, compiler_fn=compiler.transform_to_native_form) context_stack_impl.context_stack.set_default_context(context)
def test_conflicting_cardinalities_within_call(self): @computations.federated_computation([ computation_types.at_clients(tf.int32), computation_types.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): with self.assertRaisesRegex(ValueError, 'Conflicting cardinalities'): comp([five_ints, ten_ints])
def create_local_execution_context(): """Creates an XLA-based local execution context. NOTE: This context is only directly backed by an XLA executor. It does not support any intrinsics, lambda expressions, etc. Returns: An instance of `execution_context.ExecutionContext` backed by XLA executor. """ # TODO(b/175888145): Extend this into a complete local executor stack. factory = executor_stacks.local_executor_factory( support_sequence_ops=True, leaf_executor_fn=executor.XlaExecutor, local_computation_factory=compiler.XlaComputationFactory()) return execution_context.ExecutionContext(executor_fn=factory)
def create_test_execution_context(default_num_clients=0, clients_per_thread=1): """Creates an execution context that executes computations locally.""" factory = executor_stacks.local_executor_factory( default_num_clients=default_num_clients, clients_per_thread=clients_per_thread) def compiler(comp): # Compile secure_sum and secure_sum_bitwidth intrinsics to insecure # TensorFlow computations for testing purposes. replaced_intrinsic_bodies, _ = intrinsic_reductions.replace_secure_intrinsics_with_insecure_bodies( comp.to_building_block()) return computation_wrapper_instances.building_block_to_computation( replaced_intrinsic_bodies) return sync_execution_context.ExecutionContext(executor_fn=factory, compiler_fn=compiler)
def test_local_executor_multi_gpus_dataset_reduce(self, tf_device): tf_devices = tf.config.list_logical_devices(tf_device) server_tf_device = None if not tf_devices else tf_devices[0] gpu_devices = tf.config.list_logical_devices('GPU') local_executor = executor_stacks.local_executor_factory( server_tf_device=server_tf_device, client_tf_devices=gpu_devices) with executor_test_utils.install_executor(local_executor): parallel_client_run = _create_tff_parallel_clients_with_dataset_reduce() client_data = [ tf.data.Dataset.range(10), tf.data.Dataset.range(10).map(lambda x: x + 1) ] # TODO(b/159180073): merge this one into iter dataset test when the # dataset reduce function can be correctly used for GPU device. with self.assertRaisesRegex( ValueError, 'Detected dataset reduce op in multi-GPU TFF simulation.*'): parallel_client_run(client_data)
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_factory = executor_stacks.local_executor_factory(num_clients=3) tracers = [] def _tracer_fn(cardinalities): tracer = executor_test_utils.TracingExecutor( target_factory.create_executor(cardinalities)) tracers.append(tracer) return tracer service = executor_service.ExecutorService( executor_stacks.ResourceManagingExecutorFactory(_tracer_fn)) executor_pb2_grpc.add_ExecutorServicer_to_server(service, server) server.start() channel = grpc.insecure_channel('localhost:{}'.format(port)) stub = executor_pb2_grpc.ExecutorStub(channel) serialized_cards = executor_service_utils.serialize_cardinalities( {placement_literals.CLIENTS: 3}) stub.SetCardinalities( executor_pb2.SetCardinalitiesRequest(cardinalities=serialized_cards)) remote_exec = remote_executor.RemoteExecutor(channel, rpc_mode) executor = reference_resolving_executor.ReferenceResolvingExecutor( remote_exec) try: yield collections.namedtuple('_', 'executor tracers')(executor, tracers) finally: executor.close() for tracer in tracers: tracer.close() try: channel.close() except AttributeError: pass # Public gRPC channel doesn't support close() finally: server.stop(None)
def decorator(fn, *named_executors): if isinstance(fn, type): raise TypeError('Do not directly decorate classes with the executors ' 'decorator; this will cause the tests to be skipped. ' 'Decorate the member test functions instead.') if not named_executors: named_executors = [ ('local', executor_stacks.local_executor_factory()), ('sizing', executor_stacks.sizing_executor_factory()), ] @parameterized.named_parameters(*named_executors) def wrapped_fn(self, executor): """Install a particular execution context before running `fn`.""" context = execution_context.ExecutionContext(executor) with context_stack_impl.context_stack.install(context): fn(self) return wrapped_fn
def create_local_execution_context(default_num_clients: int = 0, max_fanout=100, clients_per_thread=1, server_tf_device=None, client_tf_devices=tuple(), reference_resolving_clients=False): """Creates an execution context that executes computations locally.""" factory = executor_stacks.local_executor_factory( default_num_clients=default_num_clients, max_fanout=max_fanout, clients_per_thread=clients_per_thread, server_tf_device=server_tf_device, client_tf_devices=client_tf_devices, reference_resolving_clients=reference_resolving_clients) def _compiler(comp): native_form = compiler.transform_to_native_form( comp, transform_math_to_tf=not reference_resolving_clients) return native_form return execution_context.ExecutionContext( executor_fn=factory, compiler_fn=_compiler)
def _create_concurrent_maxthread_tuples(): tuples = [] for concurrency in range(1, 5): local_ex_string = 'local_executor_{}_clients_per_thread'.format( concurrency) tf_executor_mock = ExecutorMock() ex_factory = executor_stacks.local_executor_factory( clients_per_thread=concurrency, leaf_executor_fn=tf_executor_mock) tuples.append( (local_ex_string, ex_factory, concurrency, tf_executor_mock)) sizing_ex_string = 'sizing_executor_{}_client_thread'.format( concurrency) tf_executor_mock = ExecutorMock() ex_factory = executor_stacks.sizing_executor_factory( clients_per_thread=concurrency, leaf_executor_fn=tf_executor_mock) tuples.append( (sizing_ex_string, ex_factory, concurrency, tf_executor_mock)) debug_ex_string = 'debug_executor_{}_client_thread'.format(concurrency) tf_executor_mock = ExecutorMock() ex_factory = executor_stacks.thread_debugging_executor_factory( clients_per_thread=concurrency, leaf_executor_fn=tf_executor_mock) tuples.append( (debug_ex_string, ex_factory, concurrency, tf_executor_mock)) return tuples
class IntrinsicsTest(parameterized.TestCase): def assert_type(self, value, type_string): self.assertEqual(value.type_signature.compact_representation(), type_string) def test_intrinsic_construction_raises_context_error_outside_decorator(self): @computations.tf_computation() def return_2(): return 2 with self.assertRaises(context_base.ContextError): intrinsics.federated_eval(return_2, placements.SERVER) def test_federated_broadcast_with_server_all_equal_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.SERVER)) def foo(x): val = intrinsics.federated_broadcast(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(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_eval_rand_on_clients(self): @computations.federated_computation def rand_on_clients(): @computations.tf_computation def rand(): return tf.random.normal([]) val = intrinsics.federated_eval(rand, placements.CLIENTS) self.assertIsInstance(val, value_base.Value) return val self.assert_type(rand_on_clients, '( -> {float32}@CLIENTS)') def test_federated_eval_rand_on_server(self): @computations.federated_computation def rand_on_server(): @computations.tf_computation def rand(): return tf.random.normal([]) val = intrinsics.federated_eval(rand, placements.SERVER) self.assertIsInstance(val, value_base.Value) return val self.assert_type(rand_on_server, '( -> float32@SERVER)') def test_federated_map_with_client_all_equal_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS, True)) def foo(x): val = intrinsics.federated_map( computations.tf_computation(lambda x: x > 10), x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(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): val = intrinsics.federated_map( computations.tf_computation(lambda x: x > 10), x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({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): val = intrinsics.federated_map( computations.tf_computation(lambda x: x > 10), x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(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): val = intrinsics.federated_map( computations.tf_computation(lambda x, y: x > 10, ), [x, y]) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<int32@SERVER,int32@SERVER> -> bool@SERVER)') def test_federated_map_injected_zip_fails_different_placements(self): def foo(x, y): val = intrinsics.federated_map( computations.tf_computation(lambda x, y: x > 10, ), [x, y]) self.assertIsInstance(val, value_base.Value) return val 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), x) def test_federated_sum_with_client_int(self): @computations.federated_computation( computation_types.FederatedType(tf.int32, placements.CLIENTS)) def foo(x): val = intrinsics.federated_sum(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({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): val = intrinsics.federated_zip([x, y]) self.assertIsInstance(val, value_base.Value) return val self.assert_type( foo, '(<{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): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<{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): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<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): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<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): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<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} val = intrinsics.federated_zip(a) self.assertIsInstance(val, value_base.Value) return val self.assert_type( foo, '(<{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): val = intrinsics.federated_zip([x, y]) self.assertIsInstance(val, value_base.Value) return val self.assert_type( foo, '(<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]} val = intrinsics.federated_zip(a) self.assertIsInstance(val, value_base.Value) return val self.assert_type( foo, '(<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): val = intrinsics.federated_zip([x, y]) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<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): val = intrinsics.federated_zip(arg) self.assertIsInstance(val, value_base.Value) return val self.assert_type( foo, '(<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): val = intrinsics.federated_collect(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({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): val = intrinsics.federated_mean(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({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): val = intrinsics.federated_mean(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(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): val = intrinsics.federated_mean(x, x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(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): val = intrinsics.federated_mean(x, y) self.assertIsInstance(val, value_base.Value) return val self.assert_type( foo, '(<{<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 # The operator to use during the first stage simply adds an element to the # total and updates the count. @computations.tf_computation 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 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 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): val = intrinsics.federated_aggregate(x, Accumulator(0, 0), accumulate, merge, report) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({int32}@CLIENTS -> float32@SERVER)') def test_federated_aggregate_with_federated_zero_fails(self): @computations.federated_computation() def build_federated_zero(): val = intrinsics.federated_value(0, placements.SERVER) self.assertIsInstance(val, value_base.Value) return val @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): val = intrinsics.federated_aggregate(x, build_empty_accumulator(), accumulate, merge, report) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({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) val = intrinsics.federated_reduce(x, 0, plus) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '({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): val = intrinsics.federated_sum( intrinsics.federated_map( computations.tf_computation( lambda x, y: tf.cast(tf.greater(x, y), tf.int32)), [temperatures, intrinsics.federated_broadcast(threshold)])) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(<{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) @computations.federated_computation( [computation_types.FederatedType(tf.int32, placements.CLIENTS)] * n) def foo(x): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, function_type.compact_representation()) @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) @computations.federated_computation([fed_type] * n) def foo(x): arg = {str(k): x[k] for k in range(n)} val = intrinsics.federated_zip(arg) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, named_function_type.compact_representation()) 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)) val = intrinsics.federated_zip(arg) self.assertIsInstance(val, value_base.Value) return val self.assert_type(bar, mixed_function_type.compact_representation()) @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) @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): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(baz, function_type.compact_representation()) 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): val = intrinsics.federated_apply( computations.tf_computation(lambda x: x * x), x) self.assertIsInstance(val, value_base.Value) return val self.assertLen(w, 1) self.assertIsInstance(w[0].category(), DeprecationWarning) self.assertIn('tff.federated_apply() is deprecated', str(w[0].message)) self.assert_type(foo, '(int32@SERVER -> int32@SERVER)') def test_federated_value_with_bool_on_clients(self): @computations.federated_computation(tf.bool) def foo(x): val = intrinsics.federated_value(x, placements.CLIENTS) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(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.assertIsInstance(tff_float, value_base.Value) self.assert_type(tff_float, 'float64@SERVER') intv = np.int64(0) tff_int = intrinsics.federated_value(intv, placements.SERVER) self.assertIsInstance(tff_int, value_base.Value) self.assert_type(tff_int, '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): val = intrinsics.federated_value(x, placements.SERVER) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo, '(bool -> bool@SERVER)') def test_sequence_sum(self): @computations.federated_computation( computation_types.SequenceType(tf.int32)) def foo1(x): val = intrinsics.sequence_sum(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo1, '(int32* -> int32)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.SERVER)) def foo2(x): val = intrinsics.sequence_sum(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo2, '(int32*@SERVER -> int32@SERVER)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS)) def foo3(x): val = intrinsics.sequence_sum(x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo3, '({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): val = intrinsics.sequence_map(over_threshold, x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo1, '(int32* -> bool*)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.SERVER)) def foo2(x): val = intrinsics.sequence_map(over_threshold, x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo2, '(int32*@SERVER -> bool*@SERVER)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS)) def foo3(x): val = intrinsics.sequence_map(over_threshold, x) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo3, '({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): val = intrinsics.sequence_reduce(x, 0, add_numbers) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo1, '(int32* -> int32)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.SERVER)) def foo2(x): val = intrinsics.sequence_reduce(x, 0, add_numbers) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo2, '(int32*@SERVER -> int32@SERVER)') @computations.federated_computation( computation_types.FederatedType( computation_types.SequenceType(tf.int32), placements.CLIENTS)) def foo3(x): val = intrinsics.sequence_reduce(x, 0, add_numbers) self.assertIsInstance(val, value_base.Value) return val self.assert_type(foo3, '({int32*}@CLIENTS -> {int32}@CLIENTS)') @executor_test_utils.executors( ('local', executor_stacks.local_executor_factory()),) 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): val = intrinsics.federated_zip(x) self.assertIsInstance(val, value_base.Value) return val 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 initialize_default_execution_context(): factory = executor_stacks.local_executor_factory() context = execution_context.ExecutionContext(factory) context_stack_impl.context_stack.set_default_context(context)
'The return type of next_fn must be assignable to the first parameter' ): @computations.federated_computation(tf.int32) def add_bad_result(_): return 0.0 iterative_process.IterativeProcess(initialize_fn=initialize, next_fn=add_bad_result) with self.assertRaisesRegex( TypeError, 'The return type of next_fn must be assignable to the first parameter' ): @computations.federated_computation(tf.int32) def add_bad_multi_result(_): return 0.0, 0 iterative_process.IterativeProcess(initialize_fn=initialize, next_fn=add_bad_multi_result) if __name__ == '__main__': # 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. executor = executor_stacks.local_executor_factory(num_clients=3) default_executor.set_default_executor(executor) test.main()
def _do_not_use_set_local_execution_context(): factory = executor_stacks.local_executor_factory() context = execution_context.ExecutionContext( executor_fn=factory, compiler_fn=_do_not_use_transform_to_native_form) set_default_context.set_default_context(context)
class ExecutorStacksTest(parameterized.TestCase): @parameterized.named_parameters( ('local_executor', executor_stacks.local_executor_factory), ('sizing_executor', executor_stacks.sizing_executor_factory), ('debug_executor', executor_stacks.thread_debugging_executor_factory), ) def test_construction_with_no_args(self, executor_factory_fn): executor_factory_impl = executor_factory_fn() self.assertIsInstance(executor_factory_impl, executor_stacks.ResourceManagingExecutorFactory) @parameterized.named_parameters( ('local_executor', executor_stacks.local_executor_factory), ('sizing_executor', executor_stacks.sizing_executor_factory), ) def test_construction_raises_with_max_fanout_one(self, executor_factory_fn): with self.assertRaises(ValueError): executor_factory_fn(max_fanout=1) @parameterized.named_parameters( ('local_executor_none_clients', executor_stacks.local_executor_factory()), ('sizing_executor_none_clients', executor_stacks.sizing_executor_factory()), ('local_executor_three_clients', executor_stacks.local_executor_factory(num_clients=3)), ('sizing_executor_three_clients', executor_stacks.sizing_executor_factory(num_clients=3)), ) @test_utils.skip_test_for_multi_gpu def test_execution_of_temperature_sensor_example(self, executor): comp = _temperature_sensor_example_next_fn() 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 with executor_test_utils.install_executor(executor): result = comp(temperatures, threshold) self.assertAlmostEqual(result, 8.333, places=3) @parameterized.named_parameters( ('local_executor', executor_stacks.local_executor_factory), ('sizing_executor', executor_stacks.sizing_executor_factory), ) def test_execution_with_inferred_clients_larger_than_fanout( self, executor_factory_fn): @computations.federated_computation( computation_types.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) @parameterized.named_parameters( ('local_executor_none_clients', executor_stacks.local_executor_factory()), ('sizing_executor_none_clients', executor_stacks.sizing_executor_factory()), ('debug_executor_none_clients', executor_stacks.thread_debugging_executor_factory()), ('local_executor_one_client', executor_stacks.local_executor_factory(num_clients=1)), ('sizing_executor_one_client', executor_stacks.sizing_executor_factory(num_clients=1)), ('debug_executor_one_client', executor_stacks.thread_debugging_executor_factory(num_clients=1)), ) def test_execution_of_tensorflow(self, executor): @computations.tf_computation def comp(): return tf.math.add(5, 5) with executor_test_utils.install_executor(executor): result = comp() self.assertEqual(result, 10) @parameterized.named_parameters(*_create_concurrent_maxthread_tuples()) def test_limiting_concurrency_constructs_one_eager_executor( self, ex_factory, clients_per_thread, tf_executor_mock): num_clients = 10 ex_factory.create_executor({placements.CLIENTS: num_clients}) concurrency_level = math.ceil(num_clients / clients_per_thread) args_list = tf_executor_mock.call_args_list # One for server executor, one for unplaced executor, concurrency_level for # clients. self.assertLen(args_list, concurrency_level + 2) @mock.patch( 'tensorflow_federated.python.core.impl.executors.reference_resolving_executor.ReferenceResolvingExecutor', return_value=ExecutorMock()) def test_thread_debugging_executor_constructs_exactly_one_reference_resolving_executor( self, executor_mock): executor_stacks.thread_debugging_executor_factory().create_executor( {placements.CLIENTS: 10}) executor_mock.assert_called_once() @parameterized.named_parameters( ('local_executor', executor_stacks.local_executor_factory), ('sizing_executor', executor_stacks.sizing_executor_factory), ('debug_executor', executor_stacks.thread_debugging_executor_factory), ) def test_create_executor_raises_with_wrong_cardinalities( self, executor_factory_fn): executor_factory_impl = executor_factory_fn(num_clients=5) cardinalities = { placements.SERVER: 1, None: 1, placements.CLIENTS: 1, } with self.assertRaises(ValueError, ): executor_factory_impl.create_executor(cardinalities)
class ExecutionContextIntegrationTest(parameterized.TestCase): def test_simple_no_arg_tf_computation_with_int_result(self): @computations.tf_computation def comp(): return tf.constant(10) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp() self.assertEqual(result, 10) def test_one_arg_tf_computation_with_int_param_and_result(self): @computations.tf_computation(tf.int32) def comp(x): return tf.add(x, 10) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp(3) self.assertEqual(result, 13) def test_three_arg_tf_computation_with_int_params_and_result(self): @computations.tf_computation(tf.int32, tf.int32, tf.int32) def comp(x, y, z): return tf.multiply(tf.add(x, y), z) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp(3, 4, 5) self.assertEqual(result, 35) def test_tf_computation_with_dataset_params_and_int_result(self): @computations.tf_computation(computation_types.SequenceType(tf.int32)) def comp(ds): return ds.reduce(np.int32(0), lambda x, y: x + y) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): ds = tf.data.Dataset.range(10).map(lambda x: tf.cast(x, tf.int32)) result = comp(ds) self.assertEqual(result, 45) def test_tf_computation_with_structured_result(self): @computations.tf_computation def comp(): return collections.OrderedDict([ ('a', tf.constant(10)), ('b', tf.constant(20)), ]) executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): result = comp() self.assertIsInstance(result, collections.OrderedDict) self.assertDictEqual(result, {'a': 10, 'b': 20}) @parameterized.named_parameters( ('local_executor_none_clients', executor_stacks.local_executor_factory()), ('local_executor_three_clients', executor_stacks.local_executor_factory(num_clients=3)), ) def test_with_temperature_sensor_example(self, executor): @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( computation_types.at_clients( computation_types.SequenceType(tf.float32)), computation_types.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)) with executor_test_utils.install_executor(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) def test_changing_cardinalities_across_calls(self): @computations.federated_computation( computation_types.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_conflicting_cardinalities_within_call(self): @computations.federated_computation([ computation_types.at_clients(tf.int32), computation_types.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): with self.assertRaisesRegex(ValueError, 'Conflicting cardinalities'): comp([five_ints, ten_ints]) def test_tuple_argument_can_accept_unnamed_elements(self): @computations.tf_computation(tf.int32, tf.int32) def foo(x, y): return x + y executor = executor_stacks.local_executor_factory() with executor_test_utils.install_executor(executor): # pylint:disable=no-value-for-parameter result = foo(structure.Struct([(None, 2), (None, 3)])) # pylint:enable=no-value-for-parameter self.assertEqual(result, 5)