def create_dummy_intrinsic_def_federated_map_all_equal():
value = intrinsic_defs.FEDERATED_MAP_ALL_EQUAL
type_signature = computation_types.FunctionType([
type_factory.unary_op(tf.float32),
computation_types.at_clients(tf.float32, all_equal=True),
], computation_types.at_clients(tf.float32, all_equal=True))
return value, type_signature
def create_dummy_intrinsic_def_federated_map():
value = intrinsic_defs.FEDERATED_MAP
type_signature = computation_types.FunctionType([
type_factory.unary_op(tf.float32),
type_factory.at_clients(tf.float32),
], type_factory.at_clients(tf.float32))
return value, type_signature
def create_dummy_intrinsic_def_federated_apply():
value = intrinsic_defs.FEDERATED_APPLY
type_signature = computation_types.FunctionType([
type_factory.unary_op(tf.float32),
computation_types.at_server(tf.float32),
], computation_types.at_server(tf.float32))
return value, type_signature
def create_lambda_identity(
type_spec: computation_types.Type) -> pb.Computation:
"""Returns a lambda computation representing an identity function.
Has the type signature:
(T -> T)
Args:
type_spec: A `computation_types.Type`.
Returns:
An instance of `pb.Computation`.
"""
type_signature = type_factory.unary_op(type_spec)
result = pb.Computation(type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='a'))
fn = pb.Lambda(parameter_name='a', result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
return pb.Computation(
type=type_serialization.serialize_type(type_signature),
**{'lambda': fn}) # pytype: disable=wrong-keyword-args
def test_returns_computation_sequence(self):
type_signature = computation_types.SequenceType(tf.int32)
proto = computation_factory.create_lambda_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)
def test_returns_computation_tuple_unnamed(self):
type_signature = computation_types.StructType([tf.int32, tf.float32])
proto = computation_factory.create_lambda_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)
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)
async def compute_federated_aggregate(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
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._target_executors))
identity_report = tensorflow_computation_factory.create_identity(
zero_type)
identity_report_type = type_factory.unary_op(zero_type)
aggr_type = computation_types.FunctionType(
computation_types.StructType([
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._executor.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 = await asyncio.gather(
ex.create_value(aggr_comp, aggr_type),
ex.create_struct([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._target_executors, val)])
parent_vals = await asyncio.gather(
*[self._server_executor.create_value(v, zero_type) for v in vals])
parent_merge, parent_report = await asyncio.gather(
self._server_executor.create_value(merge, merge_type),
self._server_executor.create_value(report, report_type))
merge_result = parent_vals[0]
for next_val in parent_vals[1:]:
merge_result = await self._server_executor.create_call(
parent_merge, await
self._server_executor.create_struct([merge_result, next_val]))
return FederatedComposingStrategyValue(
await self._server_executor.create_call(parent_report,
merge_result),
type_factory.at_server(report_type.result))
async def _compute_intrinsic_federated_sum(self, arg):
type_utils.check_federated_type(
arg.type_signature, placement=placement_literals.CLIENTS)
zero, plus, identity = tuple(await asyncio.gather(*[
executor_utils.embed_tf_scalar_constant(self, arg.type_signature.member,
0),
executor_utils.embed_tf_binary_operator(self, arg.type_signature.member,
tf.add),
self.create_value(
computation_factory.create_lambda_identity(
arg.type_signature.member),
type_factory.unary_op(arg.type_signature.member))
]))
aggregate_args = await self.create_tuple([arg, zero, plus, plus, identity])
return await self._compute_intrinsic_federated_aggregate(aggregate_args)
async def compute_federated_sum(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
type_analysis.check_federated_type(
arg.type_signature, placement=placement_literals.CLIENTS)
zero, plus, identity = await asyncio.gather(
executor_utils.embed_tf_scalar_constant(self._executor,
arg.type_signature.member,
0),
executor_utils.embed_tf_binary_operator(self._executor,
arg.type_signature.member,
tf.add),
self._executor.create_value(
tensorflow_computation_factory.create_identity(
arg.type_signature.member),
type_factory.unary_op(arg.type_signature.member)))
aggregate_args = await self._executor.create_struct(
[arg, zero, plus, plus, identity])
return await self.compute_federated_aggregate(aggregate_args)
def test_unary_op(self): self.assertEqual(str(type_factory.unary_op(tf.bool)), '(bool -> bool)')
def test_unary_op(self):
type_spec = computation_types.TensorType(tf.bool)
actual_type = type_factory.unary_op(type_spec)
expected_type = computation_types.FunctionType(tf.bool, tf.bool)
self.assertEqual(actual_type, expected_type)
class FederatingExecutorCreateValueTest(executor_test_utils.AsyncTestCase,
parameterized.TestCase):
# pyformat: disable
@parameterized.named_parameters([
('placement_literal',
*executor_test_utils.create_dummy_placement_literal()),
('computation_call',
*executor_test_utils.create_dummy_computation_call()),
('computation_intrinsic',
*executor_test_utils.create_dummy_computation_intrinsic()),
('computation_lambda',
*executor_test_utils.create_dummy_computation_lambda_empty()),
('computation_placement',
*executor_test_utils.create_dummy_computation_placement()),
('computation_selection',
*executor_test_utils.create_dummy_computation_selection()),
('computation_tensorflow',
*executor_test_utils.create_dummy_computation_tensorflow_empty()),
('computation_tuple',
*executor_test_utils.create_dummy_computation_tuple()),
('federated_type_at_clients',
*executor_test_utils.create_dummy_value_at_clients()),
('federated_type_at_clients_all_equal',
*executor_test_utils.create_dummy_value_at_clients_all_equal()),
('federated_type_at_server',
*executor_test_utils.create_dummy_value_at_server()),
('unplaced_type', *executor_test_utils.create_dummy_value_unplaced()),
] + get_named_parameters_for_supported_intrinsics())
# pyformat: enable
def test_returns_value_with_value_and_type(self, value, type_signature):
executor = create_test_executor()
result = self.run_sync(executor.create_value(value, type_signature))
self.assertIsInstance(result,
federating_executor.FederatingExecutorValue)
self.assertEqual(result.type_signature.compact_representation(),
type_signature.compact_representation())
# pyformat: disable
@parameterized.named_parameters([
('placement_literal',
*executor_test_utils.create_dummy_placement_literal()),
('computation_call',
*executor_test_utils.create_dummy_computation_call()),
('computation_intrinsic',
*executor_test_utils.create_dummy_computation_intrinsic()),
('computation_lambda',
*executor_test_utils.create_dummy_computation_lambda_empty()),
('computation_placement',
*executor_test_utils.create_dummy_computation_placement()),
('computation_selection',
*executor_test_utils.create_dummy_computation_selection()),
('computation_tensorflow',
*executor_test_utils.create_dummy_computation_tensorflow_empty()),
('computation_tuple',
*executor_test_utils.create_dummy_computation_tuple()),
])
# pyformat: enable
def test_returns_value_with_value_only(self, value, type_signature):
executor = create_test_executor()
result = self.run_sync(executor.create_value(value))
self.assertIsInstance(result,
federating_executor.FederatingExecutorValue)
self.assertEqual(result.type_signature.compact_representation(),
type_signature.compact_representation())
# pyformat: disable
@parameterized.named_parameters([
('computation_intrinsic',
*executor_test_utils.create_dummy_computation_intrinsic()),
('computation_lambda',
*executor_test_utils.create_dummy_computation_lambda_empty()),
('computation_tensorflow',
*executor_test_utils.create_dummy_computation_tensorflow_empty()),
])
# pyformat: enable
def test_returns_value_with_computation_impl(self, proto, type_signature):
executor = create_test_executor()
value = computation_impl.ComputationImpl(
proto, context_stack_impl.context_stack)
result = self.run_sync(executor.create_value(value, type_signature))
self.assertIsInstance(result,
federating_executor.FederatingExecutorValue)
self.assertEqual(result.type_signature.compact_representation(),
type_signature.compact_representation())
# pyformat: disable
@parameterized.named_parameters([
('federated_type_at_clients',
*executor_test_utils.create_dummy_value_at_clients()),
('federated_type_at_clients_all_equal',
*executor_test_utils.create_dummy_value_at_clients_all_equal()),
('federated_type_at_server',
*executor_test_utils.create_dummy_value_at_server()),
('unplaced_type', *executor_test_utils.create_dummy_value_unplaced()),
] + get_named_parameters_for_supported_intrinsics())
# pyformat: enable
def test_raises_type_error_with_value_only(self, value, type_signature):
del type_signature # Unused.
executor = create_test_executor()
with self.assertRaises(TypeError):
self.run_sync(executor.create_value(value))
# pyformat: disable
@parameterized.named_parameters([
('placement_literal',
*executor_test_utils.create_dummy_placement_literal()),
('computation_call',
*executor_test_utils.create_dummy_computation_call()),
('computation_intrinsic',
*executor_test_utils.create_dummy_computation_intrinsic()),
('computation_lambda',
*executor_test_utils.create_dummy_computation_lambda_empty()),
('computation_placement',
*executor_test_utils.create_dummy_computation_placement()),
('computation_selection',
*executor_test_utils.create_dummy_computation_selection()),
('computation_tensorflow',
*executor_test_utils.create_dummy_computation_tensorflow_empty()),
('computation_tuple',
*executor_test_utils.create_dummy_computation_tuple()),
('federated_type_at_clients',
*executor_test_utils.create_dummy_value_at_clients()),
('federated_type_at_clients_all_equal',
*executor_test_utils.create_dummy_value_at_clients_all_equal()),
('federated_type_at_server',
*executor_test_utils.create_dummy_value_at_server()),
('unplaced_type', *executor_test_utils.create_dummy_value_unplaced()),
] + get_named_parameters_for_supported_intrinsics())
# pyformat: enable
def test_raises_type_error_with_value_and_bad_type(self, value,
type_signature):
del type_signature # Unused.
executor = create_test_executor()
bad_type_signature = computation_types.TensorType(tf.string)
with self.assertRaises(TypeError):
self.run_sync(executor.create_value(value, bad_type_signature))
# pyformat: disable
@parameterized.named_parameters([
('computation_reference',
*executor_test_utils.create_dummy_computation_reference()),
('function_type', lambda: 10, type_factory.unary_op(tf.int32)),
])
# pyformat: enable
def test_raises_value_error_with_value(self, value, type_signature):
executor = create_test_executor()
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_unrecognized_computation_intrinsic(self):
executor = create_test_executor()
# A `ValueError` will be raised because `create_value` can not recognize the
# following intrinsic, because it has not been added to the intrinsic
# registry.
value = pb.Computation(
type=type_serialization.serialize_type(tf.int32),
intrinsic=pb.Intrinsic(uri='unregistered_intrinsic'))
type_signature = computation_types.TensorType(tf.int32)
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_unrecognized_computation_selection(self):
executor = create_test_executor()
source, _ = executor_test_utils.create_dummy_computation_tuple()
type_signature = computation_types.NamedTupleType([])
# A `ValueError` will be raised because `create_value` can not handle the
# following `pb.Selection`, because does not set either a name or an index
# field.
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
selection=pb.Selection(source=source))
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
# pyformat: disable
@parameterized.named_parameters([
('intrinsic_def_federated_broadcast',
*executor_test_utils.create_dummy_intrinsic_def_federated_broadcast()
),
('intrinsic_def_federated_eval_at_clients', *executor_test_utils.
create_dummy_intrinsic_def_federated_eval_at_clients()),
('intrinsic_def_federated_map',
*executor_test_utils.create_dummy_intrinsic_def_federated_map()),
('intrinsic_def_federated_map_all_equal', *executor_test_utils.
create_dummy_intrinsic_def_federated_map_all_equal()),
('intrinsic_def_federated_value_at_clients', *executor_test_utils.
create_dummy_intrinsic_def_federated_value_at_clients()),
('federated_type_at_clients_all_equal',
*executor_test_utils.create_dummy_value_at_clients_all_equal()),
('federated_type_at_clients',
*executor_test_utils.create_dummy_value_at_clients())
])
# pyformat: enable
def test_raises_value_error_with_no_target_executor_clients(
self, value, type_signature):
executor = federating_executor.FederatingExecutor({
placement_literals.SERVER:
eager_tf_executor.EagerTFExecutor(),
None:
eager_tf_executor.EagerTFExecutor()
})
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
# pyformat: disable
@parameterized.named_parameters([
('intrinsic_def_federated_aggregate',
*executor_test_utils.create_dummy_intrinsic_def_federated_aggregate()
),
('intrinsic_def_federated_apply',
*executor_test_utils.create_dummy_intrinsic_def_federated_apply()),
('intrinsic_def_federated_collect',
*executor_test_utils.create_dummy_intrinsic_def_federated_collect()),
('intrinsic_def_federated_eval_at_server', *executor_test_utils.
create_dummy_intrinsic_def_federated_eval_at_server()),
('intrinsic_def_federated_mean',
*executor_test_utils.create_dummy_intrinsic_def_federated_mean()),
('intrinsic_def_federated_sum',
*executor_test_utils.create_dummy_intrinsic_def_federated_sum()),
('intrinsic_def_federated_reduce',
*executor_test_utils.create_dummy_intrinsic_def_federated_reduce()),
('intrinsic_def_federated_value_at_server', *executor_test_utils.
create_dummy_intrinsic_def_federated_value_at_server()),
('intrinsic_def_federated_weighted_mean', *executor_test_utils.
create_dummy_intrinsic_def_federated_weighted_mean()),
('intrinsic_def_federated_zip_at_server', *executor_test_utils.
create_dummy_intrinsic_def_federated_zip_at_server()),
('federated_type_at_server',
*executor_test_utils.create_dummy_value_at_server()),
])
# pyformat: enable
def test_raises_value_error_with_no_target_executor_server(
self, value, type_signature):
executor = federating_executor.FederatingExecutor({
placement_literals.CLIENTS:
eager_tf_executor.EagerTFExecutor(),
None:
eager_tf_executor.EagerTFExecutor()
})
value, type_signature = executor_test_utils.create_dummy_value_at_server(
)
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_no_target_executor_unplaced(self):
executor = federating_executor.FederatingExecutor({
placement_literals.SERVER:
eager_tf_executor.EagerTFExecutor(),
placement_literals.CLIENTS:
eager_tf_executor.EagerTFExecutor(),
})
value, type_signature = executor_test_utils.create_dummy_value_unplaced(
)
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_unexpected_federated_type_at_clients(
self):
executor = create_test_executor()
value = [10, 20]
type_signature = type_factory.at_clients(tf.int32)
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
def test_raises_value_error_with_unexpected_federated_type_at_clients_all_equal(
self):
executor = create_test_executor()
value = [10] * 3
type_signature = type_factory.at_clients(tf.int32, all_equal=True)
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
