def test_with_federated_map_and_broadcast(self):
eager_ex = eager_executor.EagerExecutor()
federated_ex = federated_executor.FederatedExecutor({
None:
eager_ex,
placement_literals.SERVER:
eager_ex,
placement_literals.CLIENTS: [eager_ex for _ in range(3)]
})
ex = lambda_executor.LambdaExecutor(federated_ex)
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(
type_constructors.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(add_one,
intrinsics.federated_broadcast(x))
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(
ex.create_value(10, type_constructors.at_server(tf.int32)))
v3 = loop.run_until_complete(ex.create_call(v1, v2))
result = loop.run_until_complete(v3.compute())
self.assertCountEqual([x.numpy() for x in result], [11, 11, 11])
def test_with_federated_apply(self):
eager_ex = eager_executor.EagerExecutor()
federated_ex = federated_executor.FederatedExecutor({
None:
eager_ex,
placement_literals.SERVER:
eager_ex
})
ex = lambda_executor.LambdaExecutor(federated_ex)
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(
type_constructors.at_server(tf.int32))
def comp(x):
return intrinsics.federated_apply(add_one, x)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(
ex.create_value(10, type_constructors.at_server(tf.int32)))
v3 = loop.run_until_complete(ex.create_call(v1, v2))
result = loop.run_until_complete(v3.compute())
self.assertEqual(result.numpy(), 11)
def test_executor_create_value_with_valid_intrinsic_def(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor()
val = loop.run_until_complete(
ex.create_value(
intrinsic_defs.FEDERATED_APPLY,
computation_types.FunctionType([
type_constructors.unary_op(tf.int32),
type_constructors.at_server(tf.int32)
], type_constructors.at_server(tf.int32))))
self.assertIsInstance(val, federated_executor.FederatedExecutorValue)
self.assertEqual(str(val.type_signature),
'(<(int32 -> int32),int32@SERVER> -> int32@SERVER)')
self.assertIs(val.internal_representation,
intrinsic_defs.FEDERATED_APPLY)
def test_with_removal_of_identity_mapping(self):
@computations.federated_computation(
type_constructors.at_server(tf.int32))
def comp(x):
return intrinsics.federated_apply(_identity, x)
def transformation_fn(x):
x, _ = transformations.remove_mapped_or_applied_identity(x)
return x
self.assertEqual(_test_create_value(comp, transformation_fn),
'(FEDERATED_arg -> FEDERATED_arg)')
def test_executor_create_value_with_server_int(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor()
val = loop.run_until_complete(
ex.create_value(10, type_constructors.at_server(tf.int32)))
self.assertIsInstance(val, federated_executor.FederatedExecutorValue)
self.assertEqual(str(val.type_signature), 'int32@SERVER')
self.assertIsInstance(val.internal_representation, list)
self.assertLen(val.internal_representation, 1)
self.assertIsInstance(val.internal_representation[0],
eager_executor.EagerValue)
self.assertEqual(
val.internal_representation[0].internal_representation.numpy(), 10)
def test_with_inlining_of_blocks(self):
@computations.federated_computation(type_constructors.at_server(tf.int32))
def comp(x):
return intrinsics.federated_zip([x, x])
# TODO(b/134543154): Slide in something more powerful so that this test
# doesn't break when the implementation changes; for now, this will do.
def transformation_fn(x):
x, _ = transformations.remove_mapped_or_applied_identity(x)
x, _ = transformations.inline_block_locals(x)
x, _ = transformations.replace_selection_from_tuple_with_element(x)
return x
self.assertIn('federated_zip_at_server(<FEDERATED_arg,FEDERATED_arg>)',
_test_create_value(comp, transformation_fn))
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_constructors.at_clients(
computation_types.SequenceType(tf.float32)),
type_constructors.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))
num_clients = 3
set_default_executor.set_default_executor(
executor_stacks.create_local_executor(num_clients))
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_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_federated_mean_with_floats(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor(4)
v1 = loop.run_until_complete(
ex.create_value([1.0, 2.0, 3.0, 4.0],
type_constructors.at_clients(tf.float32)))
self.assertEqual(str(v1.type_signature), '{float32}@CLIENTS')
v2 = loop.run_until_complete(
ex.create_value(
intrinsic_defs.FEDERATED_MEAN,
computation_types.FunctionType(
type_constructors.at_clients(tf.float32),
type_constructors.at_server(tf.float32))))
self.assertEqual(str(v2.type_signature),
'({float32}@CLIENTS -> float32@SERVER)')
v3 = loop.run_until_complete(ex.create_call(v2, v1))
self.assertEqual(str(v3.type_signature), 'float32@SERVER')
result = loop.run_until_complete(v3.compute())
self.assertEqual(result.numpy(), 2.5)
def test_federated_weighted_mean_with_floats(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor(num_clients=4, use_lambda_executor=True)
v1 = loop.run_until_complete(
ex.create_value([1.0, 2.0, 3.0, 4.0],
type_constructors.at_clients(tf.float32)))
self.assertEqual(str(v1.type_signature), '{float32}@CLIENTS')
v2 = loop.run_until_complete(
ex.create_value([5.0, 10.0, 3.0, 2.0],
type_constructors.at_clients(tf.float32)))
self.assertEqual(str(v2.type_signature), '{float32}@CLIENTS')
v3 = loop.run_until_complete(
ex.create_tuple(
anonymous_tuple.AnonymousTuple([(None, v1), (None, v2)])))
self.assertEqual(str(v3.type_signature),
'<{float32}@CLIENTS,{float32}@CLIENTS>')
v4 = loop.run_until_complete(
ex.create_value(
intrinsic_defs.FEDERATED_WEIGHTED_MEAN,
computation_types.FunctionType([
type_constructors.at_clients(tf.float32),
type_constructors.at_clients(tf.float32)
], type_constructors.at_server(tf.float32))))
self.assertEqual(
str(v4.type_signature),
'(<{float32}@CLIENTS,{float32}@CLIENTS> -> float32@SERVER)')
v5 = loop.run_until_complete(ex.create_call(v4, v3))
self.assertEqual(str(v5.type_signature), 'float32@SERVER')
result = loop.run_until_complete(v5.compute())
self.assertAlmostEqual(result.numpy(), 2.1, places=3)
#
# Actual implementations might vary.
#
# Type signature: <{T}@CLIENTS,U,(<U,T>->U),(<U,U>->U),(U->R)> -> R@SERVER
FEDERATED_AGGREGATE = IntrinsicDef(
'FEDERATED_AGGREGATE', 'federated_aggregate',
computation_types.FunctionType(parameter=[
type_constructors.at_clients(computation_types.AbstractType('T')),
computation_types.AbstractType('U'),
type_constructors.reduction_op(computation_types.AbstractType('U'),
computation_types.AbstractType('T')),
type_constructors.binary_op(computation_types.AbstractType('U')),
computation_types.FunctionType(computation_types.AbstractType('U'),
computation_types.AbstractType('R'))
],
result=type_constructors.at_server(
computation_types.AbstractType('R'))))
# Applies a given function to a value on the server.
#
# Type signature: <(T->U),T@SERVER> -> U@SERVER
FEDERATED_APPLY = IntrinsicDef(
'FEDERATED_APPLY', 'federated_apply',
computation_types.FunctionType(parameter=[
computation_types.FunctionType(computation_types.AbstractType('T'),
computation_types.AbstractType('U')),
type_constructors.at_server(computation_types.AbstractType('T')),
],
result=type_constructors.at_server(
computation_types.AbstractType('U'))))
# Broadcasts a server item to all clients.
def test_at_server(self):
self.assertEqual(str(type_constructors.at_server(tf.bool)),
'bool@SERVER')
