def test_raises_value_error_with_function_type(self):
executor = create_test_executor(num_clients=3)
value = None
type_signature = type_factory.unary_op(tf.int32)
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(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),
type_factory.at_server(tf.float32),
], type_factory.at_server(tf.float32))
return value, type_signature
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),
type_factory.at_clients(tf.float32, all_equal=True),
], type_factory.at_clients(tf.float32, all_equal=True))
return value, type_signature
def create_lambda_identity(type_spec) -> pb.Computation:
"""Returns a lambda computation representing an identity function.
Has the type signature:
(T -> T)
Args:
type_spec: A type convertible to instance of `computation_types.Type` via
`computation_types.to_type`.
Returns:
An instance of `pb.Computation`.
"""
type_spec = computation_types.to_type(type_spec)
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(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)
def test_returns_computation_tuple_named(self):
type_signature = computation_types.NamedTupleType([('a', tf.int32),
('b', 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_executor_create_value_with_valid_intrinsic_def(self):
val = _produce_test_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 test_returns_computation_sequence(self):
type_signature = computation_types.SequenceType(tf.int32)
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)
expected_value = [10] * 3
actual_value = test_utils.run_tensorflow(proto, expected_value)
self.assertEqual(actual_value, expected_value)
def test_returns_computation_tuple_named(self):
type_signature = [('a', tf.int32), ('b', tf.float32)]
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)
expected_value = anonymous_tuple.AnonymousTuple([('a', 10),
('b', 10.0)])
actual_value = test_utils.run_tensorflow(proto, expected_value)
self.assertEqual(actual_value, expected_value)
def _create_lambda_identity_comp(type_spec):
py_typecheck.check_type(type_spec, computation_types.Type)
return pb.Computation(
**{
'type':
type_serialization.serialize_type(type_factory.unary_op(
type_spec)),
'lambda':
pb.Lambda(parameter_name='x',
result=pb.Computation(
type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='x')))
})
def _create_lambda_identity_comp(type_spec):
"""Returns a `pb.Computation` representing an identity function."""
py_typecheck.check_type(type_spec, computation_types.Type)
type_signature = type_serialization.serialize_type(
type_factory.unary_op(type_spec))
result = pb.Computation(type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='x'))
fn = pb.Lambda(parameter_name='x', 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_signature, **{'lambda': fn}) # pytype: disable=wrong-keyword-args
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))
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(
_create_lambda_identity_comp(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)
def create_dummy_computation_lambda_identity():
"""Returns a lambda computation and type `(float32 -> float32)`."""
type_signature = type_factory.unary_op(tf.float32)
result = pb.Computation(
type=type_serialization.serialize_type(tf.float32),
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
value = pb.Computation(
type=type_serialization.serialize_type(type_signature), **{'lambda': fn}) # pytype: disable=wrong-keyword-args
return value, type_signature
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_identity_lambda_computation(type_spec=tf.int32):
"""Returns a `pb.Computation` representing an identity lambda.
The type signature of this `pb.Computation` is:
(int32 -> int32)
Args:
type_spec: A type signature.
Returns:
A `pb.Computation`.
"""
type_signature = type_serialization.serialize_type(
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_signature, **{'lambda': fn}) # pytype: disable=wrong-keyword-args
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_placement',
*executor_test_utils.create_dummy_computation_placement()),
('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_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_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_raises_type_error_with_value_and_bad_type_skipped(
self, value, type_signature):
del type_signature # Unused.
self.skipTest(
'TODO(b/152449402): `FederatingExecutor.create_value` method should '
'fail if it is passed a computation and an incompatible type.')
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))
def test_unary_op(self):
self.assertEqual(str(type_factory.unary_op(tf.bool)), '(bool -> bool)')
def create_dummy_computation_lambda():
value = executor_test_utils.create_dummy_identity_lambda_computation()
type_signature = type_factory.unary_op(tf.int32)
return value, type_signature
def create_dummy_computation_impl():
proto = executor_test_utils.create_dummy_identity_lambda_computation()
value = computation_impl.ComputationImpl(proto,
context_stack_impl.context_stack)
type_signature = type_factory.unary_op(tf.int32)
return value, type_signature
