def test_fit_argument(self):
old_arg = reference_executor.ComputedValue(
anonymous_tuple.AnonymousTuple([('A', 10)]),
[('A', type_constructors.at_clients(tf.int32, all_equal=True))])
new_arg = reference_executor.fit_argument(
old_arg, [('A', type_constructors.at_clients(tf.int32))],
reference_executor.ComputationContext(
cardinalities={placements.CLIENTS: 3}))
self.assertEqual(str(new_arg.type_signature), '<A={int32}@CLIENTS>')
self.assertEqual(new_arg.value.A, [10, 10, 10])
def test_with_federated_value_as_a_non_py_list(self, val):
loop = asyncio.get_event_loop()
ex = _make_test_executor(num_clients=4)
v = loop.run_until_complete(
ex.create_value(val, type_constructors.at_clients(tf.int32)))
self.assertEqual(str(v.type_signature), '{int32}@CLIENTS')
result = tf.nest.map_structure(lambda x: x.numpy(),
loop.run_until_complete(v.compute()))
self.assertCountEqual(result, [1, 2, 3, 4])
def test_executor_create_value_with_all_equal_client_int(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor(3)
val = loop.run_until_complete(
ex.create_value(
10, type_constructors.at_clients(tf.int32, all_equal=True)))
self.assertIsInstance(val, federated_executor.FederatedExecutorValue)
self.assertEqual(str(val.type_signature), 'int32@CLIENTS')
self.assertIsInstance(val.internal_representation, list)
self.assertLen(val.internal_representation, 3)
for v in val.internal_representation:
self.assertIsInstance(v, eager_executor.EagerValue)
self.assertEqual(v.internal_representation.numpy(), 10)
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)
async def _compute_intrinsic_federated_broadcast(self, arg):
py_typecheck.check_type(arg.type_signature, computation_types.FederatedType)
py_typecheck.check_type(arg.internal_representation, list)
if not arg.type_signature.all_equal:
raise ValueError('Cannot broadcast a non all_equal value.')
if len(arg.internal_representation) != 1:
raise ValueError(
'Cannot broadcast a with a non-singleton representation.')
val = await arg.internal_representation[0].compute()
return FederatedExecutorValue(
await asyncio.gather(*[
c.create_value(val, arg.type_signature.member)
for c in self._target_executors[placement_literals.CLIENTS]
]),
type_constructors.at_clients(arg.type_signature.member, all_equal=True))
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()
# intrinsics defined above, as follows.
#
# @federated_computation
# def federated_aggregate(x, zero, accumulate, merge, report):
# a = generic_partial_reduce(x, zero, accumulate, INTERMEDIATE_AGGREGATORS)
# b = generic_reduce(a, zero, merge, SERVER)
# c = generic_map(report, b)
# return c
#
# 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',
def test_at_clients(self):
self.assertEqual(str(type_constructors.at_clients(tf.bool)),
'{bool}@CLIENTS')
