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