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 = lambda_executor.LambdaExecutor(remote_exec)
set_default_executor.set_default_executor(
executor_factory.ExecutorFactoryImpl(lambda _: executor))
try:
yield collections.namedtuple('_', 'executor tracer')(executor, tracer)
finally:
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.local_executor_factory(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):
"""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_executor_factory_raises_with_wrong_cardinalities(self):
ex_factory = executor_stacks.local_executor_factory(num_clients=5)
cardinalities = {
placement_literals.SERVER: 1,
None: 1,
placement_literals.CLIENTS: 1
}
with self.assertRaisesRegex(ValueError,
'construct an executor with 1 clients'):
ex_factory.create_executor(cardinalities)
@computations.tf_computation
def foo():
return tf.constant(10)
self.assertEqual(foo(), 10)
set_default_executor.set_default_executor()
with self.assertRaisesRegex(
TypeError,
'The return type of next_fn should match 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 should match 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.
set_default_executor.set_default_executor(
executor_stacks.local_executor_factory(num_clients=3))
test.main()
def test_runs_tf_unspecified_clients(self):
executor_test_utils.test_runs_tf(
self,
executor_stacks.local_executor_factory().create_executor({}))
def test_runs_tf(self):
executor_test_utils.test_runs_tf(
self,
executor_stacks.local_executor_factory(1).create_executor({}))
def test_raises_with_max_fanout_1(self):
with self.assertRaises(ValueError):
executor_stacks.local_executor_factory(2, 1)
def test_with_no_args(self):
set_default_executor.set_default_executor(
executor_stacks.local_executor_factory())
class FederatedComputationsTest(parameterized.TestCase, tf.test.TestCase):
@core_test.executors
def test_raises_value_error_none_result(self):
with self.assertRaisesRegex(ValueError, 'must return some non-`None`'):
@computations.federated_computation(None)
def _():
return None
@core_test.executors
def test_computation_with_no_args_returns_value(self):
@computations.federated_computation
def foo():
return 10
self.assertEqual(foo.type_signature.compact_representation(),
'( -> int32)')
self.assertEqual(foo(), 10)
# TODO(b/131363314): The reference executor should support generating and
# returning infinite datasets
@core_test.executors(
('local', executor_stacks.local_executor_factory(num_clients=1)), )
def test_computation_called_once_is_invoked_once(self):
@computations.tf_computation
def get_random():
return tf.random.normal([])
@computations.federated_computation
def same_random_number_twice():
value = get_random()
return value, value
num1, num2 = same_random_number_twice()
self.assertEqual(num1, num2)
@core_test.executors
def test_computation_typical_usage_as_decorator_with_unlabeled_type(self):
@computations.federated_computation(
computation_types.FunctionType(tf.int32, tf.int32),
tf.int32,
)
def foo(f, x):
assert isinstance(f, value_base.Value)
assert isinstance(x, value_base.Value)
assert str(f.type_signature) == '(int32 -> int32)'
assert str(x.type_signature) == 'int32'
result_value = f(f(x))
assert isinstance(result_value, value_base.Value)
assert str(result_value.type_signature) == 'int32'
return result_value
self.assertEqual(str(foo.type_signature),
'(<(int32 -> int32),int32> -> int32)')
@computations.tf_computation(tf.int32)
def third_power(x):
return x**3
self.assertEqual(foo(third_power, 10), int(1e9))
self.assertEqual(foo(third_power, 1), 1)
@core_test.executors
def test_computation_typical_usage_as_decorator_with_labeled_type(self):
@computations.federated_computation(
('f', computation_types.FunctionType(tf.int32, tf.int32)),
('x', tf.int32),
)
def foo(f, x):
return f(f(x))
@computations.tf_computation(tf.int32)
def square(x):
return x**2
@computations.tf_computation(tf.int32, tf.int32)
def square_drop_y(x, y): # pylint: disable=unused-argument
return x * x
self.assertEqual(str(foo.type_signature),
'(<f=(int32 -> int32),x=int32> -> int32)')
self.assertEqual(foo(square, 10), int(1e4))
self.assertEqual(square_drop_y(square_drop_y(10, 5), 100), int(1e4))
self.assertEqual(square_drop_y(square_drop_y(10, 100), 5), int(1e4))
with self.assertRaisesRegex(
TypeError,
r'(is not assignable from source type)|' # Reference executor
'(Expected a value of type .*, found .*)' # Local executor
):
foo(square_drop_y, 10)
class TensorFlowComputationsWithDatasetsTest(parameterized.TestCase):
# TODO(b/122081673): Support tf.Dataset serialization in tf2_computation.
@core_test.executors
def test_with_tf_datasets(self):
@computations.tf_computation(computation_types.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)')
@computations.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.local_executor_factory(1)), )
def test_consume_infinite_tf_dataset(self):
@computations.tf_computation(computation_types.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.local_executor_factory(1)), )
def test_produce_and_consume_infinite_tf_dataset(self):
@computations.tf_computation(computation_types.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)
@computations.tf_computation
def produce():
# Produce an infinite dataset.
return tf.data.Dataset.range(10).repeat()
self.assertEqual(consume(produce()), 45)
@core_test.executors
def test_with_sequence_of_pairs(self):
pairs = tf.data.Dataset.from_tensor_slices(
(list(range(5)), list(range(5, 10))))
@computations.tf_computation
def process_pairs(ds):
return ds.reduce(0, lambda state, pair: state + pair[0] + pair[1])
self.assertEqual(process_pairs(pairs), 45)
@core_test.executors
def test_tf_comp_with_sequence_inputs_and_outputs_does_not_fail(self):
@computations.tf_computation(computation_types.SequenceType(tf.int32))
def _(x):
return x
@core_test.executors
def test_with_four_element_dataset_pipeline(self):
@computations.tf_computation
def comp1():
return tf.data.Dataset.range(5)
@computations.tf_computation(computation_types.SequenceType(tf.int64))
def comp2(ds):
return ds.map(lambda x: tf.cast(x + 1, tf.float32))
@computations.tf_computation(computation_types.SequenceType(tf.float32)
)
def comp3(ds):
return ds.repeat(5)
@computations.tf_computation(computation_types.SequenceType(tf.float32)
)
def comp4(ds):
return ds.reduce(0.0, lambda x, y: x + y)
@computations.tf_computation
def comp5():
return comp4(comp3(comp2(comp1())))
self.assertEqual(comp5(), 75.0)
def _test_ctx(num_clients=None):
return execution_context.ExecutionContext(
executor_stacks.local_executor_factory(num_clients))
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.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):
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 _make_default_context():
return execution_context.ExecutionContext(
executor_stacks.local_executor_factory())
