def test_with_federated_map(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_factory.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(add_one, x)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(
ex.create_value(10, type_factory.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_with_federated_map_and_broadcast(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
federated_ex = federating_executor.FederatingExecutor({
None:
eager_ex,
placement_literals.SERVER:
eager_ex,
placement_literals.CLIENTS: [eager_ex for _ in range(3)]
})
ex = reference_resolving_executor.ReferenceResolvingExecutor(
federated_ex)
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(type_factory.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_factory.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_raises_with_closure(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
federated_ex = federating_executor.FederatingExecutor({
None:
eager_ex,
placement_literals.SERVER:
eager_ex,
})
ex = reference_resolving_executor.ReferenceResolvingExecutor(
federated_ex)
loop = asyncio.get_event_loop()
@computations.federated_computation(tf.int32,
type_factory.at_server(tf.int32))
def foo(x, y):
@computations.federated_computation(tf.int32)
def bar(z):
del z
return x
return intrinsics.federated_map(bar, y)
v1 = loop.run_until_complete(ex.create_value(foo))
v2 = loop.run_until_complete(
ex.create_value(
[0, 0], [tf.int32, type_factory.at_server(tf.int32)]))
with self.assertRaisesRegex(
RuntimeError,
'lambda passed to intrinsic contains references to captured variables'
):
loop.run_until_complete(ex.create_call(v1, v2))
def create_dummy_intrinsic_def_federated_apply():
value = intrinsic_defs.FEDERATED_APPLY
type_signature = computation_types.FunctionType([
type_factory.unary_op(tf.float32),
type_factory.at_server(tf.float32),
], type_factory.at_server(tf.float32))
return value, type_signature
async def _compute_intrinsic_federated_aggregate(self, arg):
value_type, zero_type, accumulate_type, merge_type, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
py_typecheck.check_type(arg.internal_representation,
anonymous_tuple.AnonymousTuple)
py_typecheck.check_len(arg.internal_representation, 5)
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
py_typecheck.check_len(val, len(self._child_executors))
identity_report = _create_lambda_identity_comp(zero_type)
identity_report_type = type_factory.unary_op(zero_type)
aggr_type = computation_types.FunctionType(
computation_types.NamedTupleType([
value_type, zero_type, accumulate_type, merge_type,
identity_report_type
]), type_factory.at_server(zero_type))
aggr_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_AGGREGATE, aggr_type)
zero = await (await self.create_selection(arg, index=1)).compute()
accumulate = arg.internal_representation[2]
merge = arg.internal_representation[3]
report = arg.internal_representation[4]
async def _child_fn(ex, v):
py_typecheck.check_type(v, executor_value_base.ExecutorValue)
aggr_func, aggr_args = tuple(await asyncio.gather(
ex.create_value(aggr_comp, aggr_type),
ex.create_tuple([v] + list(await asyncio.gather(
ex.create_value(zero, zero_type),
ex.create_value(accumulate, accumulate_type),
ex.create_value(merge, merge_type),
ex.create_value(identity_report, identity_report_type))))))
return await (await ex.create_call(aggr_func, aggr_args)).compute()
vals = await asyncio.gather(
*[_child_fn(c, v) for c, v in zip(self._child_executors, val)])
parent_vals = await asyncio.gather(
*[self._parent_executor.create_value(v, zero_type) for v in vals])
parent_merge, parent_report = tuple(await asyncio.gather(
self._parent_executor.create_value(merge, merge_type),
self._parent_executor.create_value(report, report_type)))
merge_result = parent_vals[0]
for next_val in parent_vals[1:]:
merge_result = await self._parent_executor.create_call(
parent_merge, await
self._parent_executor.create_tuple([merge_result, next_val]))
return CompositeValue(
await self._parent_executor.create_call(parent_report,
merge_result),
type_factory.at_server(report_type.result))
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_factory.unary_op(tf.int32),
type_factory.at_server(tf.int32)
], type_factory.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 create_dummy_intrinsic_def():
"""Returns a `intrinsic_defs.IntrinsicDef` and type."""
value = intrinsic_defs.FEDERATED_EVAL_AT_SERVER
type_signature = computation_types.FunctionType(
computation_types.FunctionType(None, tf.int32),
type_factory.at_server(tf.int32))
return value, type_signature
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(
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))
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)
class CreateIdentityTest(parameterized.TestCase):
# pyformat: disable
@parameterized.named_parameters(
('int', computation_types.TensorType(tf.int32), 10),
('unnamed_tuple',
computation_types.NamedTupleType([tf.int32, tf.float32]),
anonymous_tuple.AnonymousTuple([(None, 10), (None, 10.0)])),
('named_tuple',
computation_types.NamedTupleType([
('a', tf.int32), ('b', tf.float32)
]), anonymous_tuple.AnonymousTuple([('a', 10), ('b', 10.0)])),
('sequence', computation_types.SequenceType(tf.int32), [10] * 3),
)
# pyformat: enable
def test_returns_computation(self, type_signature, value):
proto = tensorflow_computation_factory.create_identity(type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = type_factory.unary_op(type_signature)
self.assertEqual(actual_type, expected_type)
actual_result = test_utils.run_tensorflow(proto, value)
self.assertEqual(actual_result, value)
@parameterized.named_parameters(
('none', None),
('federated_type', type_factory.at_server(tf.int32)),
)
def test_raises_type_error(self, type_signature):
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_identity(type_signature)
def test_get_size_info(self, num_clients):
@computations.federated_computation(
type_factory.at_clients(computation_types.SequenceType(tf.float32)),
type_factory.at_server(tf.float32))
def comp(temperatures, threshold):
client_data = [temperatures, intrinsics.federated_broadcast(threshold)]
result_map = intrinsics.federated_map(
count_over, intrinsics.federated_zip(client_data))
count_map = intrinsics.federated_map(count_total, temperatures)
return intrinsics.federated_mean(result_map, count_map)
factory = executor_stacks.sizing_executor_factory(num_clients=num_clients)
default_executor.set_default_executor(factory)
to_float = lambda x: tf.cast(x, tf.float32)
temperatures = [tf.data.Dataset.range(10).map(to_float)] * num_clients
threshold = 15.0
comp(temperatures, threshold)
# Each client receives a tf.float32 and uploads two tf.float32 values.
expected_broadcast_bits = num_clients * 32
expected_aggregate_bits = expected_broadcast_bits * 2
expected = ({
(('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients
}, {
(('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients * 2
}, [expected_broadcast_bits], [expected_aggregate_bits])
self.assertEqual(expected, factory.get_size_info())
def _temperature_sensor_example_next_fn():
@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))
return comp
def create_dummy_intrinsic_def_federated_secure_sum():
value = intrinsic_defs.FEDERATED_SECURE_SUM
type_signature = computation_types.FunctionType([
type_factory.at_clients(tf.float32),
tf.float32,
], type_factory.at_server(tf.float32))
return value, type_signature
async def _get_cardinalities_helper():
"""Helper function which does the actual work of fetching cardinalities."""
one_type = type_factory.at_clients(tf.int32, all_equal=True)
sum_type = computation_types.FunctionType(
type_factory.at_clients(tf.int32),
type_factory.at_server(tf.int32))
sum_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_SUM, sum_type)
async def _count_leaf_executors(ex):
"""Counts the total number of leaf executors under `ex`."""
one_fut = ex.create_value(1, one_type)
sum_comp_fut = ex.create_value(sum_comp, sum_type)
one_val, sum_comp_val = tuple(await asyncio.gather(
one_fut, sum_comp_fut))
sum_result = await (await
ex.create_call(sum_comp_val,
one_val)).compute()
if isinstance(sum_result, tf.Tensor):
return sum_result.numpy()
else:
return sum_result
return await asyncio.gather(
*[_count_leaf_executors(c) for c in self._child_executors])
def test_returns_value_with_federated_type_server(self):
value = [eager_tf_executor.EagerValue(10.0, None, tf.float32)]
type_signature = type_factory.at_server(tf.float32)
value = federating_executor.FederatingExecutorValue(value, type_signature)
result = self.run_sync(value.compute())
self.assertEqual(result, 10.0)
def create_dummy_intrinsic_def_federated_reduce():
value = intrinsic_defs.FEDERATED_REDUCE
type_signature = computation_types.FunctionType([
type_factory.at_clients(tf.float32),
tf.float32,
type_factory.reduction_op(tf.float32, tf.float32),
], type_factory.at_server(tf.float32))
return value, type_signature
class CreateBinaryOperatorTest(parameterized.TestCase):
# pyformat: disable
@parameterized.named_parameters(
('add_int', tf.math.add, computation_types.TensorType(
tf.int32), [1, 2], 3),
('add_float', tf.math.add, computation_types.TensorType(
tf.float32), [1.0, 2.25], 3.25),
('add_unnamed_tuple', tf.math.add,
computation_types.NamedTupleType([tf.int32, tf.float32]), [
[1, 1.0], [2, 2.25]
], anonymous_tuple.AnonymousTuple([(None, 3), (None, 3.25)])),
('add_named_tuple', tf.math.add,
computation_types.NamedTupleType([
('a', tf.int32), ('b', tf.float32)
]), [[1, 1.0], [2, 2.25]
], anonymous_tuple.AnonymousTuple([('a', 3), ('b', 3.25)])),
('multiply_int', tf.math.multiply,
computation_types.TensorType(tf.int32), [2, 2], 4),
('multiply_float', tf.math.multiply,
computation_types.TensorType(tf.float32), [2.0, 2.25], 4.5),
('divide_int', tf.math.divide, computation_types.TensorType(
tf.int32), [4, 2], 2.0),
('divide_float', tf.math.divide,
computation_types.TensorType(tf.float32), [4.0, 2.0], 2.0),
('divide_inf', tf.math.divide, computation_types.TensorType(
tf.int32), [1, 0], np.inf),
)
# pyformat: enable
def test_returns_computation(self, operator, type_signature, operands,
expected_result):
proto = tensorflow_computation_factory.create_binary_operator(
operator, type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
self.assertIsInstance(actual_type, computation_types.FunctionType)
# Note: It is only useful to test the parameter type; the result type
# depends on the `operator` used, not the implemenation
# `create_binary_operator`.
expected_parameter_type = computation_types.NamedTupleType(
[type_signature, type_signature])
self.assertEqual(actual_type.parameter, expected_parameter_type)
actual_result = test_utils.run_tensorflow(proto, operands)
self.assertEqual(actual_result, expected_result)
@parameterized.named_parameters(
('non_callable_operator', 1, computation_types.TensorType(tf.int32)),
('none_type', tf.math.add, None),
('federated_type', tf.math.add, type_factory.at_server(tf.int32)),
('sequence_type', tf.math.add, computation_types.SequenceType(
tf.int32)),
)
def test_raises_type_error(self, operator, type_signature):
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_binary_operator(
operator, type_signature)
def _federated_mean(self, arg):
type_utils.check_federated_type(arg.type_signature, None,
placements.CLIENTS, False)
py_typecheck.check_type(arg.value, list)
server_sum = self._federated_sum(arg)
unplaced_avg = multiply_by_scalar(
ComputedValue(server_sum.value, server_sum.type_signature.member),
1.0 / float(len(arg.value)))
return ComputedValue(unplaced_avg.value,
type_factory.at_server(unplaced_avg.type_signature))
def test_serialize_deserialize_federated_at_server(self):
x = 10
x_type = type_factory.at_server(tf.int32)
value_proto, value_type = executor_service_utils.serialize_value(
x, x_type)
self.assertIsInstance(value_proto, executor_pb2.Value)
self.assertEqual(str(value_type), 'int32@SERVER')
y, type_spec = executor_service_utils.deserialize_value(value_proto)
self.assertEqual(str(type_spec), str(x_type))
self.assertEqual(y, 10)
def test_with_removal_of_identity_mapping(self):
@computations.federated_computation(type_factory.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(_identity, x)
def transformation_fn(x):
x, _ = tree_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_factory.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 _federated_zip_at_server(self, arg):
py_typecheck.check_type(arg.type_signature,
computation_types.NamedTupleType)
for idx in range(len(arg.type_signature)):
type_utils.check_federated_type(arg.type_signature[idx], None,
placements.SERVER, True)
return ComputedValue(
arg.value,
type_factory.at_server(
computation_types.NamedTupleType([
(k, v.member) if k else v.member
for k, v in anonymous_tuple.iter_elements(arg.type_signature)
])))
def test_returns_value_with_unplaced_type_and_server(self, executor):
value, type_signature = executor_test_utils.create_dummy_value_unplaced(
)
value = self.run_sync(executor.create_value(value, type_signature))
result = self.run_sync(
executor_utils.compute_intrinsic_federated_value(
executor, value, placement_literals.SERVER))
self.assertIsInstance(result, executor_value_base.ExecutorValue)
expected_type = type_factory.at_server(type_signature)
self.assertEqual(result.type_signature.compact_representation(),
expected_type.compact_representation())
actual_result = self.run_sync(result.compute())
self.assertEqual(actual_result, 10.0)
def test_with_inlining_of_blocks(self):
@computations.federated_computation(type_factory.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))
class CreateConstantTest(parameterized.TestCase):
# pyformat: disable
@parameterized.named_parameters(
('int', 10, computation_types.TensorType(tf.int32, [3]), [10] * 3),
('float', 10.0, computation_types.TensorType(tf.float32,
[3]), [10.0] * 3),
('unnamed_tuple', 10, computation_types.NamedTupleType(
[tf.int32] * 3), anonymous_tuple.AnonymousTuple([(None, 10)] * 3)),
('named_tuple', 10,
computation_types.NamedTupleType([
('a', tf.int32), ('b', tf.int32), ('c', tf.int32)
]), anonymous_tuple.AnonymousTuple([('a', 10), ('b', 10),
('c', 10)])),
('nested_tuple', 10,
computation_types.NamedTupleType([[tf.int32] * 3] * 3),
anonymous_tuple.AnonymousTuple(
[(None, anonymous_tuple.AnonymousTuple([(None, 10)] * 3))] * 3)),
)
# pyformat: enable
def test_returns_computation(self, value, type_signature, expected_result):
proto = tensorflow_computation_factory.create_constant(
value, type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = computation_types.FunctionType(None, type_signature)
self.assertEqual(actual_type, expected_type)
actual_result = test_utils.run_tensorflow(proto)
if isinstance(expected_result, list):
self.assertCountEqual(actual_result, expected_result)
else:
self.assertEqual(actual_result, expected_result)
@parameterized.named_parameters(
('non_scalar_value', np.zeros(
[1]), computation_types.TensorType(tf.int32)),
('none_type', 10, None),
('federated_type', 10, type_factory.at_server(tf.int32)),
('bad_type', 10.0, computation_types.TensorType(tf.int32)),
)
def test_raises_type_error(self, value, type_signature):
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_constant(
value, type_signature)
