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(num_clients=3)
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 proto(self):
return pb.Computation(
type=type_serialization.serialize_type(self.type_signature),
block=pb.Block(
**{
'local': [
pb.Block.Local(name=k, value=v.proto)
for k, v in self._locals
],
'result': self._result.proto
}))
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 proto(self):
elements = []
for k, v in anonymous_tuple.to_elements(self):
if k is not None:
element = pb.Tuple.Element(name=k, value=v.proto)
else:
element = pb.Tuple.Element(value=v.proto)
elements.append(element)
return pb.Computation(
type=type_serialization.serialize_type(self.type_signature),
tuple=pb.Tuple(element=elements))
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))
def test_data_proto_tensor(self):
ex = data_executor.DataExecutor(
eager_tf_executor.EagerTFExecutor(),
TestDataBackend(self, 'foo://bar', 10, tf.int32))
proto = pb.Computation(data=pb.Data(uri='foo://bar'),
type=type_serialization.serialize_type(
computation_types.TensorType(tf.int32)))
val = self._loop.run_until_complete(ex.create_value(proto))
self.assertIsInstance(val, eager_tf_executor.EagerValue)
self.assertEqual(str(val.type_signature), 'int32')
self.assertEqual(self._loop.run_until_complete(val.compute()), 10)
ex.close()
def create_dummy_computation_tuple():
"""Returns a tuple computation and type."""
names = ['a', 'b', 'c']
element_value, element_type = create_dummy_computation_tensorflow_constant(
)
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 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))
def create_lambda_empty_struct() -> pb.Computation:
"""Returns a lambda computation returning an empty struct.
Has the type signature:
( -> <>)
Returns:
An instance of `pb.Computation`.
"""
result_type = computation_types.StructType([])
type_signature = computation_types.FunctionType(None, result_type)
result = pb.Computation(
type=type_serialization.serialize_type(result_type),
struct=pb.Struct(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
return pb.Computation(
type=type_serialization.serialize_type(type_signature), **{'lambda': fn}) # pytype: disable=wrong-keyword-args
def test_raises_on_xla(self):
function_type = computation_types.FunctionType(
computation_types.TensorType(tf.int32),
computation_types.TensorType(tf.int32))
empty_xla_computation_proto = computation_pb2.Computation(
type=type_serialization.serialize_type(function_type),
xla=computation_pb2.Xla())
compiled_comp = building_blocks.CompiledComputation(
proto=empty_xla_computation_proto)
with self.assertRaises(compiler.XlaToTensorFlowError):
compiler.compile_local_computation_to_tensorflow(compiled_comp)
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 main(_: Sequence[str]) -> None:
def ex_fn(device: tf.config.LogicalDevice) -> tff.framework.DataExecutor:
# In order to de-reference data uri's bundled in TFF computations, a
# DataExecutor must exist in the runtime context to process those uri's and
# return the underlying data. We can wrap an EagerTFExecutor (which handles
# TF operations) with a DataExecutor instance defined with a DataBackend
# object.
return tff.framework.DataExecutor(
tff.framework.EagerTFExecutor(device),
data_backend=NumpyArrDataBackend())
# Executor factory used by the runtime context to spawn executors to run TFF
# computations.
factory = tff.framework.local_executor_factory(leaf_executor_fn=ex_fn)
# Context in which to execute the following computation.
ctx = tff.framework.ExecutionContext(executor_fn=factory)
tff.framework.set_default_context(ctx)
# Type of the data returned by the DataBackend.
element_type = tff.types.TensorType(tf.int32)
element_type_proto = tff.framework.serialize_type(element_type)
# We construct a list of uri's as our references to the dataset.
uris = [f'uri://{i}' for i in range(3)]
# The uris are embedded in TFF computation protos so they can be processed by
# TFF executors.
arguments = [
pb.Computation(data=pb.Data(uri=uri), type=element_type_proto)
for uri in uris
]
# The embedded uris are passed to a DataDescriptor which recognizes the
# underlying dataset as federated and allows combining it with a federated
# computation.
data_handle = tff.framework.DataDescriptor(
None, arguments, tff.FederatedType(element_type, tff.CLIENTS),
len(arguments))
# Federated computation that sums the values in the arrays.
@tff.federated_computation(tff.types.FederatedType(element_type, tff.CLIENTS))
def foo(x):
@tff.tf_computation(element_type)
def local_sum(nums):
return tf.math.reduce_sum(nums)
return tff.federated_sum(tff.federated_map(local_sum, x))
# Should print 18.
print(foo(data_handle))
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 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
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_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 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_invoke_raises_value_error_with_federated_computation(self):
bogus_proto = pb.Computation(type=type_serialization.serialize_type(
computation_types.to_type(
computation_types.FunctionType(tf.int32, tf.int32))),
reference=pb.Reference(name='boogledy'))
non_tf_computation = computation_impl.ComputationImpl(
bogus_proto, context_stack_impl.context_stack)
context = tensorflow_computation_context.TensorFlowComputationContext(
tf.compat.v1.get_default_graph())
with self.assertRaisesRegex(
ValueError, 'Can only invoke TensorFlow in the body of '
'a TensorFlow computation'):
context.invoke(non_tf_computation, None)
def wrap_graph_parameter_as_tuple(comp, name=None):
"""Wraps the parameter of `comp` in a tuple binding.
`wrap_graph_parameter_as_tuple` is intended as a preprocessing step
to `pad_graph_inputs_to_match_type`, so that `pad_graph_inputs_to_match_type`
can
make the assumption that its argument `comp` always has a tuple binding,
instead of dealing with the possibility of an unwrapped tensor or sequence
binding.
Args:
comp: Instance of `computation_building_blocks.CompiledComputation` whose
parameter we wish to wrap in a tuple binding.
name: Optional string argument, the name to assign to the element type in
the constructed tuple. Defaults to `None`.
Returns:
A transformed version of comp representing exactly the same computation,
but accepting a tuple containing one element--the parameter of `comp`.
Raises:
TypeError: If `comp` is not a
`computation_building_blocks.CompiledComputation`.
"""
py_typecheck.check_type(comp,
computation_building_blocks.CompiledComputation)
if name is not None:
py_typecheck.check_type(name, six.string_types)
proto = comp.proto
proto_type = type_serialization.deserialize_type(proto.type)
parameter_binding = [proto.tensorflow.parameter]
parameter_type_list = [(name, proto_type.parameter)]
new_parameter_binding = pb.TensorFlow.Binding(
tuple=pb.TensorFlow.NamedTupleBinding(element=parameter_binding))
new_function_type = computation_types.FunctionType(parameter_type_list,
proto_type.result)
serialized_type = type_serialization.serialize_type(new_function_type)
input_padded_proto = pb.Computation(
type=serialized_type,
tensorflow=pb.TensorFlow(graph_def=proto.tensorflow.graph_def,
initialize_op=proto.tensorflow.initialize_op,
parameter=new_parameter_binding,
result=proto.tensorflow.result))
return computation_building_blocks.CompiledComputation(input_padded_proto)
def create_dummy_computation_tensorflow_empty():
"""Returns a tensorflow computation and type `( -> <>)`."""
with tf.Graph().as_default() as graph:
result_type, result_binding = tensorflow_utils.capture_result_from_graph(
[], 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 disable_grappler_for_partitioned_calls(proto):
"""Disables grappler for `PartitionedCall` and `StatefulPartitionedCall` nodes in the graph.
TensorFlow serializes a `ConfigProto` into `PartitionedCall` and
`StatefulPartitionedCall` the `config_proto` `attr` of graph nodes. This
overrides any session config that might disable runtime grappler. The disable
grappler for these nodes as well, this function overwrites the serialized
configproto, setting the `disable_meta_optimizer` field to `True.
Args:
proto: Instance of `computation_pb2.Computation` with the `tensorflow` field
populated.
Returns:
A transformed instance of `computation_pb2.Computation` with a `tensorflow`
field.
"""
py_typecheck.check_type(proto, computation_pb2.Computation)
computation_oneof = proto.WhichOneof('computation')
if computation_oneof != 'tensorflow':
raise TypeError('`prune_tensorflow_proto` only accepts `Computation` '
'protos of the "tensorflow" variety; you have passed '
'one of variety {}.'.format(computation_oneof))
original_tf = proto.tensorflow
graph_def = serialization_utils.unpack_graph_def(original_tf.graph_def)
all_nodes = itertools.chain(
graph_def.node, *[f.node_def for f in graph_def.library.function])
for node in all_nodes:
if node.op not in CALL_OPS:
continue
attr_str = node.attr.get('config_proto')
if attr_str is None:
config_proto = tf.compat.v1.ConfigProto()
else:
config_proto = tf.compat.v1.ConfigProto.FromString(attr_str.s)
config_proto.graph_options.rewrite_options.disable_meta_optimizer = True
attr_str.s = config_proto.SerializeToString(deterministic=True)
tf_block = computation_pb2.TensorFlow(
graph_def=serialization_utils.pack_graph_def(graph_def),
initialize_op=original_tf.initialize_op
if original_tf.initialize_op else None,
parameter=original_tf.parameter
if original_tf.HasField('parameter') else None,
result=original_tf.result)
new_proto = computation_pb2.Computation(type=proto.type,
tensorflow=tf_block)
return new_proto
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_get_wrapped_function_from_comp_raises_with_incorrect_binding(self):
with tf.Graph().as_default() as graph:
var = tf.Variable(initial_value=0.0, name='var1', import_scope='')
assign_op = var.assign_add(tf.constant(1.0))
tf.add(1.0, assign_op)
result_binding = pb.TensorFlow.Binding(
tensor=pb.TensorFlow.TensorBinding(tensor_name='Invalid'))
comp = pb.Computation(
tensorflow=pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(graph.as_graph_def()),
result=result_binding))
with self.assertRaises(TypeError):
wrapped_fn = eager_tf_executor._get_wrapped_function_from_comp(
comp, must_pin_function_to_cpu=False, param_type=None, device=None)
wrapped_fn()
def create_xla_tff_computation(xla_computation, type_spec):
"""Creates an XLA TFF computation.
Args:
xla_computation: An instance of `xla_client.XlaComputation`.
type_spec: The TFF type of the computation to be constructed.
Returns:
An instance of `pb.Computation`.
"""
py_typecheck.check_type(xla_computation, xla_client.XlaComputation)
py_typecheck.check_type(type_spec, computation_types.FunctionType)
return pb.Computation(
type=type_serialization.serialize_type(type_spec),
xla=pb.Xla(hlo_module=pack_xla_computation(xla_computation),
parameter=_make_xla_binding_for_type(type_spec.parameter),
result=_make_xla_binding_for_type(type_spec.result)))
def test_data_proto_dataset(self):
type_spec = computation_types.SequenceType(tf.int64)
ex = data_executor.DataExecutor(
eager_tf_executor.EagerTFExecutor(),
TestDataBackend(self, 'foo://bar', tf.data.Dataset.range(3),
type_spec))
proto = pb.Computation(
data=pb.Data(uri='foo://bar'),
type=type_serialization.serialize_type(type_spec))
val = self._loop.run_until_complete(ex.create_value(proto))
self.assertIsInstance(val, eager_tf_executor.EagerValue)
self.assertEqual(str(val.type_signature), 'int64*')
self.assertCountEqual([
x.numpy()
for x in iter(self._loop.run_until_complete(val.compute()))
], [0, 1, 2])
ex.close()
def prune_tensorflow_proto(proto):
"""Extracts subgraph from `proto` preserving parameter, result and initialize.
Args:
proto: Instance of `pb.Computation` of the `tensorflow` variety whose
`graphdef` attribute we wish to prune of extraneous ops.
Returns:
A transformed instance of `pb.Computation` of the `tensorflow` variety,
whose `graphdef` attribute contains only ops which can reach the
parameter or result bindings, or initialize op.
"""
py_typecheck.check_type(proto, pb.Computation)
computation_oneof = proto.WhichOneof('computation')
if computation_oneof != 'tensorflow':
raise TypeError(
'`prune_tensorflow_proto` only accepts `Computation` '
'protos of the \'tensorflow\' variety; you have passed '
'one of variety {}.'.format(computation_oneof))
if proto.tensorflow.parameter.WhichOneof('binding'):
parameter_tensor_names = graph_utils.extract_tensor_names_from_binding(
proto.tensorflow.parameter)
parameter_names = [
':'.join(x.split(':')[:-1]) for x in parameter_tensor_names
]
else:
parameter_names = []
return_tensor_names = graph_utils.extract_tensor_names_from_binding(
proto.tensorflow.result)
return_names = [':'.join(x.split(':')[:-1]) for x in return_tensor_names]
graph_def = serialization_utils.unpack_graph_def(
proto.tensorflow.graph_def)
init_op_name = proto.tensorflow.initialize_op
names_to_preserve = parameter_names + return_names
if init_op_name:
names_to_preserve.append(init_op_name)
subgraph_def = tf.compat.v1.graph_util.extract_sub_graph(
graph_def, names_to_preserve)
tf_block = pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(subgraph_def),
initialize_op=proto.tensorflow.initialize_op,
parameter=proto.tensorflow.parameter,
result=proto.tensorflow.result)
pruned_proto = pb.Computation(type=proto.type, tensorflow=tf_block)
return pruned_proto
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))
