def test_executor_create_value_with_intrinsic_as_pb_computation(self):
val = _produce_test_value(
pb.Computation(intrinsic=pb.Intrinsic(uri='generic_zero'),
type=type_serialization.serialize_type(tf.int32)))
self.assertIsInstance(val, federated_executor.FederatedExecutorValue)
self.assertEqual(str(val.type_signature), 'int32')
self.assertIs(val.internal_representation, intrinsic_defs.GENERIC_ZERO)
def create_dummy_computation_intrinsic():
"""Returns a intrinsic computation and type."""
intrinsic_def, type_signature = create_dummy_intrinsic_def()
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
intrinsic=pb.Intrinsic(uri=intrinsic_def.uri))
return value, type_signature
def create_whimsy_computation_intrinsic():
"""Returns a intrinsic computation and type."""
intrinsic_def, type_signature = create_whimsy_intrinsic_def_federated_eval_at_server(
)
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
intrinsic=pb.Intrinsic(uri=intrinsic_def.uri))
return value, type_signature
def test_executor_call_unsupported_intrinsic(self):
dummy_intrinsic = intrinsic_defs.IntrinsicDef(
'DUMMY_INTRINSIC', 'dummy_intrinsic',
computation_types.AbstractType('T'))
comp = pb.Computation(intrinsic=pb.Intrinsic(uri='dummy_intrinsic'),
type=type_serialization.serialize_type(tf.int32))
with self.assertRaises(NotImplementedError):
_run_test_comp(comp, num_clients=3)
def test_executor_create_value_with_intrinsic_as_pb_computation(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor()
val = loop.run_until_complete(
ex.create_value(
pb.Computation(intrinsic=pb.Intrinsic(uri='generic_zero'),
type=type_serialization.serialize_type(
tf.int32))))
self.assertIsInstance(val, federated_executor.FederatedExecutorValue)
self.assertEqual(str(val.type_signature), 'int32')
self.assertIs(val.internal_representation, intrinsic_defs.GENERIC_ZERO)
def test_raises_not_implemented_error_with_unimplemented_intrinsic(self):
executor = create_test_executor()
dummy_intrinsic = intrinsic_defs.IntrinsicDef(
'DUMMY_INTRINSIC', 'dummy_intrinsic',
computation_types.AbstractType('T'))
comp = pb.Computation(intrinsic=pb.Intrinsic(uri='dummy_intrinsic'),
type=type_serialization.serialize_type(tf.int32))
comp = self.run_sync(executor.create_value(comp))
with self.assertRaises(NotImplementedError):
self.run_sync(executor.create_call(comp))
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 create_intrinsic_comp(intrinsic_def, type_spec):
"""Creates an intrinsic `pb.Computation`.
Args:
intrinsic_def: An instance of `intrinsic_defs.IntrinsicDef`.
type_spec: The concrete type of the intrinsic (`computation_types.Type`).
Returns:
An instance of `pb.Computation` that represents the intrinsics.
"""
py_typecheck.check_type(intrinsic_def, intrinsic_defs.IntrinsicDef)
py_typecheck.check_type(type_spec, computation_types.Type)
return pb.Computation(type=type_serialization.serialize_type(type_spec),
intrinsic=pb.Intrinsic(uri=intrinsic_def.uri))
def test_executor_call_unsupported_intrinsic(self):
dummy_intrinsic = intrinsic_defs.IntrinsicDef(
'DUMMY_INTRINSIC', 'dummy_intrinsic',
computation_types.AbstractType('T'))
comp = pb.Computation(type=type_serialization.serialize_type(tf.int32),
intrinsic=pb.Intrinsic(uri='dummy_intrinsic'))
loop = asyncio.get_event_loop()
executor = composing_executor.ComposingExecutor(
_create_bottom_stack(), [_create_worker_stack() for _ in range(3)])
with self.assertRaises(NotImplementedError):
v1 = loop.run_until_complete(executor.create_value(comp, tf.int32))
loop.run_until_complete(executor.create_call(v1, None))
def test_with_type_raises_non_assignable_type(self):
int_return_type = computation_types.FunctionType(tf.int32, tf.int32)
original_comp = computation_impl.ConcreteComputation(
pb.Computation(
**{
'type': type_serialization.serialize_type(int_return_type),
'intrinsic': pb.Intrinsic(uri='whatever')
}), context_stack_impl.context_stack)
list_return_type = computation_types.FunctionType(
tf.int32,
computation_types.StructWithPythonType([(None, tf.int32)], list))
with self.assertRaises(computation_types.TypeNotAssignableError):
computation_impl.ConcreteComputation.with_type(original_comp,
list_return_type)
def test_raises_not_implemented_error_with_unimplemented_intrinsic(self):
executor = create_test_executor()
# `whimsy_intrinsic` definition is needed to allow lookup.
whimsy_intrinsic = intrinsic_defs.IntrinsicDef(
'WHIMSY_INTRINSIC', 'whimsy_intrinsic',
computation_types.AbstractType('T'))
type_signature = computation_types.TensorType(tf.int32)
comp = pb.Computation(
intrinsic=pb.Intrinsic(uri='whimsy_intrinsic'),
type=type_serialization.serialize_type(type_signature))
del whimsy_intrinsic
comp = self.run_sync(executor.create_value(comp))
with self.assertRaises(NotImplementedError):
self.run_sync(executor.create_call(comp))
def test_with_type_preserves_python_container(self):
struct_return_type = computation_types.FunctionType(
tf.int32, computation_types.StructType([(None, tf.int32)]))
original_comp = computation_impl.ConcreteComputation(
pb.Computation(
**{
'type': type_serialization.serialize_type(struct_return_type),
'intrinsic': pb.Intrinsic(uri='whatever')
}), context_stack_impl.context_stack)
list_return_type = computation_types.FunctionType(
tf.int32,
computation_types.StructWithPythonType([(None, tf.int32)], list))
fn_with_annotated_type = computation_impl.ConcreteComputation.with_type(
original_comp, list_return_type)
type_test_utils.assert_types_identical(
list_return_type, fn_with_annotated_type.type_signature)
def test_executor_call_unsupported_intrinsic(self):
dummy_intrinsic = intrinsic_defs.IntrinsicDef(
'DUMMY_INTRINSIC', 'dummy_intrinsic',
computation_types.AbstractType('T'))
type_signature = computation_types.TensorType(tf.int32)
comp = pb.Computation(
type=type_serialization.serialize_type(type_signature),
intrinsic=pb.Intrinsic(uri='dummy_intrinsic'))
loop = asyncio.get_event_loop()
factory = federated_composing_strategy.FederatedComposingStrategy.factory(
_create_bottom_stack(), [_create_worker_stack()])
executor = federating_executor.FederatingExecutor(
factory, _create_bottom_stack())
v1 = loop.run_until_complete(executor.create_value(comp))
with self.assertRaises(NotImplementedError):
loop.run_until_complete(executor.create_call(v1))
def test_executor_call_unsupported_intrinsic(self):
# `whimsy_intrinsic` definition is needed to allow successful lookup.
whimsy_intrinsic = intrinsic_defs.IntrinsicDef(
'WHIMSY_INTRINSIC', 'whimsy_intrinsic',
computation_types.AbstractType('T'))
type_signature = computation_types.TensorType(tf.int32)
comp = pb.Computation(
type=type_serialization.serialize_type(type_signature),
intrinsic=pb.Intrinsic(uri='whimsy_intrinsic'))
del whimsy_intrinsic
factory = federated_composing_strategy.FederatedComposingStrategy.factory(
_create_bottom_stack(), [_create_worker_stack()])
executor = federating_executor.FederatingExecutor(
factory, _create_bottom_stack())
v1 = asyncio.run(executor.create_value(comp))
with self.assertRaises(NotImplementedError):
asyncio.run(executor.create_call(v1))
def test_returns_value_with_intrinsic_def_federated_value_at_server_and_tuple(
self):
executor = create_test_executor(num_clients=3)
arg, arg_type = executor_test_utils.create_dummy_computation_tuple()
intrinsic_def = intrinsic_defs.FEDERATED_VALUE_AT_SERVER
comp_type = computation_types.FunctionType(arg_type,
type_factory.at_server(arg_type))
comp = pb.Computation(
type=type_serialization.serialize_type(comp_type),
intrinsic=pb.Intrinsic(uri=intrinsic_def.uri))
comp = self.run_sync(executor.create_value(comp, comp_type))
arg = self.run_sync(executor.create_value(arg, arg_type))
result = self.run_sync(executor.create_call(comp, arg))
self.assertIsInstance(result, federating_executor.FederatingExecutorValue)
self.assertEqual(result.type_signature.compact_representation(),
comp_type.result.compact_representation())
actual_result = self.run_sync(result.compute())
expected_result = [10.0] * 2
for actual_element, expected_element in zip(actual_result, expected_result):
self.assertEqual(actual_element, expected_element)
def test_something(self):
# TODO(b/113112108): Revise these tests after a more complete implementation
# is in place.
# At the moment, this should succeed, as both the computation body and the
# type are well-formed.
computation_impl.ComputationImpl(
pb.Computation(
**{
'type':
type_serialization.serialize_type(
computation_types.FunctionType(tf.int32, tf.int32)),
'intrinsic':
pb.Intrinsic(uri='whatever')
}), context_stack_impl.context_stack)
# This should fail, as the proto is not well-formed.
self.assertRaises(TypeError, computation_impl.ComputationImpl,
pb.Computation(), context_stack_impl.context_stack)
# This should fail, as "10" is not an instance of pb.Computation.
self.assertRaises(TypeError, computation_impl.ComputationImpl, 10,
context_stack_impl.context_stack)
def _make_sequence_map_value(executor, source_type, target_type):
intrinsic_spec = pb.Intrinsic(uri=intrinsic_defs.SEQUENCE_MAP.uri)
type_spec = _make_sequence_map_type(source_type, target_type)
comp_pb = pb.Computation(type=type_serialization.serialize_type(type_spec),
intrinsic=intrinsic_spec)
return _run_sync(executor.create_value(comp_pb, type_spec))
def proto(self):
return pb.Computation(type=type_serialization.serialize_type(
self.type_signature),
intrinsic=pb.Intrinsic(uri=self._uri))
def create_dummy_computation_intrinsic():
value = pb.Computation(
type=type_serialization.serialize_type(tf.int32),
intrinsic=pb.Intrinsic(uri=intrinsic_defs.GENERIC_ZERO.uri))
type_signature = computation_types.TensorType(tf.int32)
return value, type_signature
def _make_sequence_reduce_value(executor, element_type, accumulator_type):
intrinsic_spec = pb.Intrinsic(uri=intrinsic_defs.SEQUENCE_REDUCE.uri)
type_spec = _make_sequence_reduce_type(element_type, accumulator_type)
comp_pb = pb.Computation(type=type_serialization.serialize_type(type_spec),
intrinsic=intrinsic_spec)
return _run_sync(executor.create_value(comp_pb, type_spec))
