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 create_dummy_computation_tuple():
"""Returns a tuple computation and type."""
element_type = computation_types.NamedTupleType([])
element_value = pb.Computation(
type=type_serialization.serialize_type(element_type),
tuple=pb.Tuple(element=[]))
element = pb.Tuple.Element(value=element_value)
type_signature = computation_types.NamedTupleType([element_type])
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
tuple=pb.Tuple(element=[element]))
return value, type_signature
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
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_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 _serialize_deserialize_roundtrip_test(self, type_list):
"""Performs roundtrip serialization/deserialization of computation_types.
Args:
type_list: A list of instances of computation_types.Type or things
convertible to it.
"""
for t in type_list:
t1 = computation_types.to_type(t)
p1 = type_serialization.serialize_type(t1)
t2 = type_serialization.deserialize_type(p1)
p2 = type_serialization.serialize_type(t2)
self.assertEqual(repr(t1), repr(t2))
self.assertEqual(repr(p1), repr(p2))
self.assertTrue(type_utils.are_equivalent_types(t1, t2))
def create_dummy_computation_lambda_empty():
"""Returns a lambda computation and type `( -> <>)`."""
result_type = computation_types.NamedTupleType([])
type_signature = computation_types.FunctionType(None, result_type)
result = pb.Computation(
type=type_serialization.serialize_type(result_type),
tuple=pb.Tuple(element=[]))
fn = pb.Lambda(parameter_name=None, 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 create_dummy_computation_tuple():
"""Returns a tuple computation and type."""
names = ['a', 'b', 'c']
fn, fn_type = create_dummy_computation_tensorflow_constant()
element_value = pb.Computation(
type=type_serialization.serialize_type(fn_type),
call=pb.Call(function=fn))
element_type = fn_type.result
elements = [pb.Tuple.Element(name=n, value=element_value) for n in names]
type_signature = computation_types.NamedTupleType(
(n, element_type) for n in names)
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
tuple=pb.Tuple(element=elements))
return value, type_signature
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_dummy_computation_call():
function = executor_test_utils.create_dummy_empty_tensorflow_computation()
value = pb.Computation(type=type_serialization.serialize_type(
computation_types.NamedTupleType([])),
call=pb.Call(function=function))
type_signature = computation_types.NamedTupleType([])
return value, type_signature
def create_dummy_computation_selection():
element_value = executor_test_utils.create_dummy_empty_tensorflow_computation(
)
element_type = computation_types.FunctionType(
None, computation_types.NamedTupleType([]))
element = pb.Tuple.Element(value=element_value)
source = pb.Computation(type=type_serialization.serialize_type(
[element_type]),
tuple=pb.Tuple(element=[element]))
value = pb.Computation(
type=type_serialization.serialize_type(element_type),
selection=pb.Selection(source=source, index=0))
type_signature = computation_types.FunctionType(
None, computation_types.NamedTupleType([]))
return value, type_signature
def create_dummy_computation_placement():
value = pb.Computation(
type=type_serialization.serialize_type(
computation_types.PlacementType()),
placement=pb.Placement(uri=placement_literals.SERVER.uri))
type_signature = computation_types.PlacementType()
return value, type_signature
def create_replicate_input(type_spec, count: int) -> pb.Computation:
"""Returns a tensorflow computation which returns `count` clones of an input.
The returned computation has the type signature `(T -> <T, T, T, ...>)`, where
`T` is `type_spec` and the length of the result is `count`.
Args:
type_spec: A type convertible to instance of `computation_types.Type` via
`computation_types.to_type`.
count: An integer, the number of times the input is replicated.
Raises:
TypeError: If `type_spec` contains any types which cannot appear in
TensorFlow bindings or if `which` is not an integer.
"""
type_spec = computation_types.to_type(type_spec)
type_utils.check_tensorflow_compatible_type(type_spec)
py_typecheck.check_type(count, int)
with tf.Graph().as_default() as graph:
parameter_value, parameter_binding = tensorflow_utils.stamp_parameter_in_graph(
'x', type_spec, graph)
result = [parameter_value] * count
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
result, graph)
type_signature = computation_types.FunctionType(type_spec, result_type)
tensorflow = pb.TensorFlow(graph_def=serialization_utils.pack_graph_def(
graph.as_graph_def()),
parameter=parameter_binding,
result=result_binding)
return pb.Computation(
type=type_serialization.serialize_type(type_signature),
tensorflow=tensorflow)
def create_identity(type_spec) -> pb.Computation:
"""Returns a tensorflow computation representing an identity function.
The returned computation has the type signature `(T -> T)`, where `T` is
`type_spec`.
Args:
type_spec: A type convertible to instance of `computation_types.Type` via
`computation_types.to_type`.
Raises:
TypeError: If `type_spec` contains any types which cannot appear in
TensorFlow bindings.
"""
type_spec = computation_types.to_type(type_spec)
type_utils.check_tensorflow_compatible_type(type_spec)
with tf.Graph().as_default() as graph:
parameter_value, parameter_binding = tensorflow_utils.stamp_parameter_in_graph(
'x', type_spec, graph)
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
parameter_value, graph)
type_signature = computation_types.FunctionType(type_spec, result_type)
tensorflow = pb.TensorFlow(graph_def=serialization_utils.pack_graph_def(
graph.as_graph_def()),
parameter=parameter_binding,
result=result_binding)
return pb.Computation(
type=type_serialization.serialize_type(type_signature),
tensorflow=tensorflow)
def create_dummy_computation_tensorflow_tuple(value=10.0):
"""Returns a tensorflow computation and type.
`( -> <('a', T), ('b', T), ('c', T)>)`
Args:
value: An optional integer value.
"""
with tf.Graph().as_default() as graph:
names = ['a', 'b', 'c']
result = anonymous_tuple.AnonymousTuple(
(n, tf.constant(value)) for n in names)
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
result, graph)
type_signature = computation_types.FunctionType(None, result_type)
tensorflow = pb.TensorFlow(graph_def=serialization_utils.pack_graph_def(
graph.as_graph_def()),
parameter=None,
result=result_binding)
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
tensorflow=tensorflow)
return value, type_signature
def test_gets_some_explicit_some_none_placement(self):
with tf.Graph().as_default() as g:
with tf.device('/cpu:0'):
a = tf.constant(0)
b = tf.constant(1)
c = a + b
_, result_binding = tensorflow_utils.capture_result_from_graph(c, g)
packed_graph_def = serialization_utils.pack_graph_def(g.as_graph_def())
function_type = computation_types.FunctionType(None, tf.int32)
proto = pb.Computation(
type=type_serialization.serialize_type(function_type),
tensorflow=pb.TensorFlow(graph_def=packed_graph_def,
parameter=None,
result=result_binding))
building_block = building_blocks.ComputationBuildingBlock.from_proto(
proto)
device_placements = building_block_analysis.get_device_placement_in(
building_block)
all_device_placements = list(device_placements.keys())
self.assertLen(all_device_placements, 2)
if all_device_placements[0]:
self.assertIn('CPU', all_device_placements[0])
self.assertEqual('', all_device_placements[1])
else:
self.assertIn('CPU', all_device_placements[1])
self.assertEqual('', all_device_placements[0])
self.assertGreater(device_placements[all_device_placements[0]], 0)
self.assertGreater(device_placements[all_device_placements[1]], 0)
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_tensorflow_add():
"""Returns a tensorflow computation and type.
`(<float32,float32> -> float32)`
"""
type_spec = tf.float32
with tf.Graph().as_default() as graph:
parameter_1_value, parameter_1_binding = tensorflow_utils.stamp_parameter_in_graph(
'x', type_spec, graph)
parameter_2_value, parameter_2_binding = tensorflow_utils.stamp_parameter_in_graph(
'y', type_spec, graph)
result_value = tf.add(parameter_1_value, parameter_2_value)
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
result_value, graph)
parameter_type = computation_types.NamedTupleType([type_spec, type_spec])
type_signature = computation_types.FunctionType(parameter_type, result_type)
tuple_binding = pb.TensorFlow.NamedTupleBinding(
element=[parameter_1_binding, parameter_2_binding])
parameter_binding = pb.TensorFlow.Binding(tuple=tuple_binding)
tensorflow = pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(graph.as_graph_def()),
parameter=parameter_binding,
result=result_binding)
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
tensorflow=tensorflow)
return value, type_signature
def create_dummy_computation_reference():
"""Returns a reference computation and type."""
type_signature = computation_types.TensorType(tf.float32)
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
reference=pb.Reference(name='a'))
return value, type_signature
def create_dummy_computation_placement():
"""Returns a placement computation and type."""
placement_literal, type_signature = create_dummy_placement_literal()
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
placement=pb.Placement(uri=placement_literal.uri))
return value, type_signature
def create_dummy_computation_selection():
"""Returns a selection computation and type."""
source, source_type = create_dummy_computation_tuple()
type_signature = source_type[0]
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
selection=pb.Selection(source=source, index=0))
return value, type_signature
def create_dummy_computation_call():
"""Returns a call computation and type."""
fn, fn_type = create_dummy_computation_tensorflow_constant()
type_signature = fn_type.result
value = pb.Computation(
type=type_serialization.serialize_type(type_signature),
call=pb.Call(function=fn))
return value, type_signature
def _pack_noarg_graph(graph_def, return_type, result_binding):
packed_graph_def = serialization_utils.pack_graph_def(graph_def)
function_type = computation_types.FunctionType(None, return_type)
proto = pb.Computation(
type=type_serialization.serialize_type(function_type),
tensorflow=pb.TensorFlow(
graph_def=packed_graph_def, parameter=None, result=result_binding))
building_block = building_blocks.ComputationBuildingBlock.from_proto(proto)
return building_block
def test_serialize_type_with_federated_bool(self):
actual_proto = type_serialization.serialize_type(
computation_types.FederatedType(tf.bool, placements.CLIENTS, True))
expected_proto = pb.Type(federated=pb.FederatedType(
placement=pb.PlacementSpec(value=pb.Placement(
uri=placements.CLIENTS.uri)),
all_equal=True,
member=_create_scalar_tensor_type(tf.bool)))
self.assertEqual(actual_proto, expected_proto)
def test_executor_create_value_with_unbound_reference(self):
loop = asyncio.get_event_loop()
ex = _make_test_executor()
with self.assertRaises(ValueError):
loop.run_until_complete(
ex.create_value(
pb.Computation(reference=pb.Reference(name='a'),
type=type_serialization.serialize_type(
tf.int32))))
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_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_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_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(num_clients=3)
element_value = executor_test_utils.create_dummy_empty_tensorflow_computation(
)
element_type = computation_types.FunctionType(
None, computation_types.NamedTupleType([]))
element = pb.Tuple.Element(value=element_value)
source = pb.Computation(type=type_serialization.serialize_type(
[element_type]),
tuple=pb.Tuple(element=[element]))
# 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(element_type),
selection=pb.Selection(source=source))
type_signature = computation_types.FunctionType(
None, computation_types.NamedTupleType([]))
with self.assertRaises(ValueError):
self.run_sync(executor.create_value(value, type_signature))
