def test_reduces_unnecessary_ops(self): proto = _create_proto_with_unnecessary_op() comp = building_blocks.CompiledComputation(proto) reduced_proto = proto_transformations.prune_tensorflow_proto(proto) reduced_comp = building_blocks.CompiledComputation(reduced_proto) ops_before = building_block_analysis.count_tensorflow_ops_in(comp) ops_after = building_block_analysis.count_tensorflow_ops_in(reduced_comp) self.assertLess(ops_after, ops_before)
def test_does_not_reduce_no_unnecessary_ops(self):
comp = building_block_factory.create_compiled_identity(tf.int32)
pruned = building_blocks.CompiledComputation(
proto_transformations.prune_tensorflow_proto(comp.proto))
ops_before = building_block_analysis.count_tensorflow_ops_in(comp)
ops_after = building_block_analysis.count_tensorflow_ops_in(pruned)
self.assertEqual(ops_before, ops_after)
def _wrap_computation_as_value(proto: pb.Computation) -> Value:
"""Wraps the given computation as a `tff.Value`."""
py_typecheck.check_type(proto, pb.Computation)
compiled = building_blocks.CompiledComputation(proto)
call = building_blocks.Call(compiled)
ref = _bind_computation_to_reference(call,
'wrapping a computation as a value')
return Value(ref)
def transform(self, comp):
if not self.should_transform(comp):
return comp, False
py_typecheck.check_type(comp, building_blocks.CompiledComputation)
new_comp_proto = tensorflow_computation_transformations.disable_grappler_for_partitioned_calls(
comp.proto)
return building_blocks.CompiledComputation(
new_comp_proto, type_signature=comp.type_signature), True
def _transform(building_block):
nonlocal counter
if building_block.is_compiled_computation():
new_name = str(counter)
counter += 1
return building_blocks.CompiledComputation(
proto=building_block.proto, name=new_name), True
return building_block, False
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_reduces_unnecessary_ops(self):
def bad_fn(x):
_ = tf.constant(0)
return x
comp = _create_compiled_computation(bad_fn, tf.int32)
ops_before = building_block_analysis.count_tensorflow_ops_in(comp)
reduced_proto = proto_transformations.prune_tensorflow_proto(
comp.proto)
reduced_comp = building_blocks.CompiledComputation(reduced_proto)
ops_after = building_block_analysis.count_tensorflow_ops_in(
reduced_comp)
self.assertLess(ops_after, ops_before)
def test_prune_does_not_change_exeuction(self):
def bad_fn(x):
_ = tf.constant(0)
return x
comp = _create_compiled_computation(bad_fn, tf.int32)
reduced_proto = proto_transformations.prune_tensorflow_proto(
comp.proto)
reduced_comp = building_blocks.CompiledComputation(reduced_proto)
orig_executable = computation_wrapper_instances.building_block_to_computation(
comp)
reduced_executable = computation_wrapper_instances.building_block_to_computation(
reduced_comp)
for k in range(5):
self.assertEqual(orig_executable(k), reduced_executable(k))
def test_get_curried(self):
operand_type = computation_types.TensorType(tf.int32)
computation_proto, type_signature = tensorflow_computation_factory.create_binary_operator(
tf.add, operand_type, operand_type)
building_block = building_blocks.CompiledComputation(
proto=computation_proto,
name='test',
type_signature=type_signature)
add_numbers = value_impl.Value(building_block)
curried = value_utils.get_curried(add_numbers)
self.assertEqual(curried.type_signature.compact_representation(),
'(int32 -> (int32 -> int32))')
self.assertEqual(curried.comp.compact_representation(),
'(arg0 -> (arg1 -> comp#test(<arg0,arg1>)))')
def _wrap_constant_as_value(const, context_stack):
"""Wraps the given Python constant as a `tff.Value`.
Args:
const: Python constant to be converted to TFF value. Anything convertible to
Tensor via `tf.constant` can be passed in.
context_stack: The context stack to use.
Returns:
An instance of `value_base.Value`.
"""
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
tf_comp, _ = tensorflow_serialization.serialize_py_fn_as_tf_computation(
lambda: tf.constant(const), None, context_stack)
compiled_comp = building_blocks.CompiledComputation(tf_comp)
called_comp = building_blocks.Call(compiled_comp)
return ValueImpl(called_comp, context_stack)
def transform(self, comp):
if not self.should_transform(comp):
return comp, False
py_typecheck.check_type(comp, building_blocks.CompiledComputation)
new_tf_proto = computation_pb2.TensorFlow()
new_tf_proto.CopyFrom(comp.proto.tensorflow)
# Important: we must also serialize the type_signature because TFF might
# produce (<> -> <>) or (<> -> <<>>) functions, which both could be
# represented as the same graph with a single NoOp node. This can occur
# particularly in MapReduceForm compiltion for secure_sum intrinsics over
# empty structures.
hash_value = hash(
(comp.type_signature, comp.proto.tensorflow.graph_def.value))
new_tf_proto.cache_key.id = ctypes.c_uint64(hash_value).value
new_comp_proto = computation_pb2.Computation(
type=comp.proto.type, tensorflow=new_tf_proto)
return building_blocks.CompiledComputation(
new_comp_proto, type_signature=comp.type_signature), True
def _wrap_computation_as_value(
proto: pb.Computation,
context_stack: context_stack_base.ContextStack) -> value_base.Value:
"""Wraps the given computation as a `tff.Value`.
Args:
proto: A pb.Computation.
context_stack: The context stack to use.
Returns:
A `value_base.Value`.
"""
py_typecheck.check_type(proto, pb.Computation)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
compiled = building_blocks.CompiledComputation(proto)
call = building_blocks.Call(compiled)
federated_computation_context = context_stack.current
ref = federated_computation_context.bind_computation_to_reference(call)
return ValueImpl(ref, context_stack)
def _wrap_sequence_as_value(elements, element_type, context_stack):
"""Wraps `elements` as a TFF sequence with elements of type `element_type`.
Args:
elements: Python object to the wrapped as a TFF sequence value.
element_type: An instance of `Type` that determines the type of elements of
the sequence.
context_stack: The context stack to use.
Returns:
An instance of `tff.Value`.
Raises:
TypeError: If `elements` and `element_type` are of incompatible types.
"""
# TODO(b/113116813): Add support for other representations of sequences.
py_typecheck.check_type(elements, list)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
# Checks that the types of all the individual elements are compatible with the
# requested type of the sequence as a while.
for elem in elements:
elem_type = type_conversions.infer_type(elem)
if not element_type.is_assignable_from(elem_type):
raise TypeError(
'Expected all sequence elements to be {}, found {}.'.format(
element_type, elem_type))
# Defines a no-arg function that builds a `tf.data.Dataset` from the elements.
def _create_dataset_from_elements():
return tensorflow_utils.make_data_set_from_elements(
tf.compat.v1.get_default_graph(), elements, element_type)
# Wraps the dataset as a value backed by a no-argument TensorFlow computation.
tf_comp, _ = tensorflow_serialization.serialize_py_fn_as_tf_computation(
_create_dataset_from_elements, None, context_stack)
# TODO(b/159281959): Follow up and bind a reference here.
return ValueImpl(
building_blocks.Call(building_blocks.CompiledComputation(tf_comp)),
context_stack)
def optimize_tensorflow_comp(tf_computation, config_proto):
"""Applies configured optimizations to the graphdef backing a TF comp.
Args:
tf_computation: Instance of `building_blocks.CompiledComputation` backed by
TensorFlow.
config_proto: Instance of `tf.compat.v1.ConfigProto` specifying the
optimizations to apply to the graph backing this TensorFlow computation.
Returns:
A transformed version of `tf_computation`, which has had the
`tf.compat.v1.GraphDef` backing it run through Grappler with the specified
configuration.
"""
py_typecheck.check_type(tf_computation, building_blocks.CompiledComputation)
tf_proto = tf_computation.proto
graph_spec_obj = _unpack_proto_into_graph_spec(tf_proto)
optimized_graph_spec = graph_optimizations.optimize_graph_spec(
graph_spec_obj, config_proto)
graph_def = serialization_utils.pack_graph_def(optimized_graph_spec.graph_def)
original_tf = tf_proto.tensorflow
tf_result_proto = computation_pb2.TensorFlow(
graph_def=graph_def,
initialize_op=(original_tf.initialize_op
if original_tf.initialize_op else None),
session_token_tensor_name=(original_tf.session_token_tensor_name
if original_tf.session_token_tensor_name else
None),
parameter=(original_tf.parameter
if original_tf.HasField('parameter') else None),
result=original_tf.result)
optimized_proto = computation_pb2.Computation(
type=tf_proto.type, tensorflow=tf_result_proto)
return building_blocks.CompiledComputation(
optimized_proto, type_signature=tf_computation.type_signature)
def _create_compiled_computation(py_fn, arg_type):
proto, _ = tensorflow_serialization.serialize_py_fn_as_tf_computation(
py_fn, arg_type, context_stack_impl.context_stack)
return building_blocks.CompiledComputation(proto)
def to_compiled_building_block(self):
return building_blocks.CompiledComputation(
self._computation_proto, type_signature=self.type_signature)
def to_value(
arg: Any,
type_spec,
context_stack: context_stack_base.ContextStack,
) -> ValueImpl:
"""Converts the argument into an instance of `tff.Value`.
The types of non-`tff.Value` arguments that are currently convertible to
`tff.Value` include the following:
* Lists, tuples, anonymous tuples, named tuples, and dictionaries, all
of which are converted into instances of `tff.Tuple`.
* Placement literals, converted into instances of `tff.Placement`.
* Computations.
* Python constants of type `str`, `int`, `float`, `bool`
* Numpy objects inherting from `np.ndarray` or `np.generic` (the parent
of numpy scalar types)
Args:
arg: Either an instance of `tff.Value`, or an argument convertible to
`tff.Value`. The argument must not be `None`.
type_spec: An optional `computation_types.Type` or value convertible to it
by `computation_types.to_type` which specifies the desired type signature
of the resulting value. This allows for disambiguating the target type
(e.g., when two TFF types can be mapped to the same Python
representations), or `None` if none available, in which case TFF tries to
determine the type of the TFF value automatically.
context_stack: The context stack to use.
Returns:
An instance of `tff.Value` corresponding to the given `arg`, and of TFF type
matching the `type_spec` if specified (not `None`).
Raises:
TypeError: if `arg` is of an unsupported type, or of a type that does not
match `type_spec`. Raises explicit error message if TensorFlow constructs
are encountered, as TensorFlow code should be sealed away from TFF
federated context.
"""
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
if type_spec is not None:
type_spec = computation_types.to_type(type_spec)
type_utils.check_well_formed(type_spec)
if isinstance(arg, ValueImpl):
result = arg
elif isinstance(arg, building_blocks.ComputationBuildingBlock):
result = ValueImpl(arg, context_stack)
elif isinstance(arg, placement_literals.PlacementLiteral):
result = ValueImpl(building_blocks.Placement(arg), context_stack)
elif isinstance(arg, computation_base.Computation):
result = ValueImpl(
building_blocks.CompiledComputation(
computation_impl.ComputationImpl.get_proto(arg)),
context_stack)
elif type_spec is not None and isinstance(type_spec,
computation_types.SequenceType):
result = _wrap_sequence_as_value(arg, type_spec.element, context_stack)
elif isinstance(arg, anonymous_tuple.AnonymousTuple):
result = ValueImpl(
building_blocks.Tuple([
(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in anonymous_tuple.iter_elements(arg)
]), context_stack)
elif py_typecheck.is_named_tuple(arg):
result = to_value(arg._asdict(), None, context_stack) # pytype: disable=attribute-error
elif py_typecheck.is_attrs(arg):
result = to_value(
attr.asdict(arg,
dict_factory=collections.OrderedDict,
recurse=False), None, context_stack)
elif isinstance(arg, dict):
if isinstance(arg, collections.OrderedDict):
items = arg.items()
else:
items = sorted(arg.items())
value = building_blocks.Tuple([
(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in items
])
result = ValueImpl(value, context_stack)
elif isinstance(arg, (tuple, list)):
result = ValueImpl(
building_blocks.Tuple([
ValueImpl.get_comp(to_value(x, None, context_stack))
for x in arg
]), context_stack)
elif isinstance(arg, tensorflow_utils.TENSOR_REPRESENTATION_TYPES):
result = _wrap_constant_as_value(arg, context_stack)
elif isinstance(arg, (tf.Tensor, tf.Variable)):
raise TypeError(
'TensorFlow construct {} has been encountered in a federated '
'context. TFF does not support mixing TF and federated orchestration '
'code. Please wrap any TensorFlow constructs with '
'`tff.tf_computation`.'.format(arg))
elif isinstance(arg, function_utils.PolymorphicFunction):
# TODO(b/129567727) remove this case when this is no longer an error
raise TypeError(
'Polymorphic computations cannot be converted to a TFF value. Consider '
'explicitly specifying the argument types of a computation before '
'passing it to a function that requires a TFF value (such as a TFF '
'intrinsic like federated_map).')
else:
raise TypeError(
'Unable to interpret an argument of type {} as a TFF value.'.
format(py_typecheck.type_string(type(arg))))
py_typecheck.check_type(result, ValueImpl)
if (type_spec is not None and not type_utils.is_assignable_from(
type_spec, result.type_signature)):
raise TypeError(
'The supplied argument maps to TFF type {}, which is incompatible with '
'the requested type {}.'.format(result.type_signature, type_spec))
return result
async def compute_intrinsic_federated_weighted_mean(
executor: executor_base.Executor,
arg: executor_value_base.ExecutorValue,
local_computation_factory: local_computation_factory_base.
LocalComputationFactory = tensorflow_computation_factory.
TensorFlowComputationFactory()
) -> executor_value_base.ExecutorValue:
"""Computes a federated weighted mean on the given `executor`.
Args:
executor: The executor to use.
arg: The argument to embedded in `executor`.
local_computation_factory: An instance of `LocalComputationFactory` to use
to construct local computations used as parameters in certain federated
operators (such as `tff.federated_sum`, etc.). Defaults to a TensorFlow
computation factory that generates TensorFlow code.
Returns:
The result embedded in `executor`.
"""
type_analysis.check_valid_federated_weighted_mean_argument_tuple_type(
arg.type_signature)
zip1_type = computation_types.FunctionType(
computation_types.StructType([
computation_types.at_clients(arg.type_signature[0].member),
computation_types.at_clients(arg.type_signature[1].member)
]),
computation_types.at_clients(
computation_types.StructType(
[arg.type_signature[0].member, arg.type_signature[1].member])))
operand_type = zip1_type.result.member[0]
scalar_type = zip1_type.result.member[1]
multiply_comp_pb, multiply_comp_type = local_computation_factory.create_scalar_multiply_operator(
operand_type, scalar_type)
multiply_blk = building_blocks.CompiledComputation(
multiply_comp_pb, type_signature=multiply_comp_type)
map_type = computation_types.FunctionType(
computation_types.StructType(
[multiply_blk.type_signature, zip1_type.result]),
computation_types.at_clients(multiply_blk.type_signature.result))
sum1_type = computation_types.FunctionType(
computation_types.at_clients(map_type.result.member),
computation_types.at_server(map_type.result.member))
sum2_type = computation_types.FunctionType(
computation_types.at_clients(arg.type_signature[1].member),
computation_types.at_server(arg.type_signature[1].member))
zip2_type = computation_types.FunctionType(
computation_types.StructType([sum1_type.result, sum2_type.result]),
computation_types.at_server(
computation_types.StructType(
[sum1_type.result.member, sum2_type.result.member])))
divide_blk = building_block_factory.create_tensorflow_binary_operator_with_upcast(
zip2_type.result.member, tf.divide)
async def _compute_multiply_fn():
return await executor.create_value(multiply_blk.proto,
multiply_blk.type_signature)
async def _compute_multiply_arg():
zip1_comp = create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS, zip1_type)
zip_fn = await executor.create_value(zip1_comp, zip1_type)
return await executor.create_call(zip_fn, arg)
async def _compute_product_fn():
map_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_MAP,
map_type)
return await executor.create_value(map_comp, map_type)
async def _compute_product_arg():
multiply_fn, multiply_arg = await asyncio.gather(
_compute_multiply_fn(), _compute_multiply_arg())
return await executor.create_struct((multiply_fn, multiply_arg))
async def _compute_products():
product_fn, product_arg = await asyncio.gather(_compute_product_fn(),
_compute_product_arg())
return await executor.create_call(product_fn, product_arg)
async def _compute_total_weight():
sum2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM,
sum2_type)
sum2_fn, sum2_arg = await asyncio.gather(
executor.create_value(sum2_comp, sum2_type),
executor.create_selection(arg, 1))
return await executor.create_call(sum2_fn, sum2_arg)
async def _compute_sum_of_products():
sum1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM,
sum1_type)
sum1_fn, products = await asyncio.gather(
executor.create_value(sum1_comp, sum1_type), _compute_products())
return await executor.create_call(sum1_fn, products)
async def _compute_zip2_fn():
zip2_comp = create_intrinsic_comp(
intrinsic_defs.FEDERATED_ZIP_AT_SERVER, zip2_type)
return await executor.create_value(zip2_comp, zip2_type)
async def _compute_zip2_arg():
sum_of_products, total_weight = await asyncio.gather(
_compute_sum_of_products(), _compute_total_weight())
return await executor.create_struct([sum_of_products, total_weight])
async def _compute_divide_fn():
return await executor.create_value(divide_blk.proto,
divide_blk.type_signature)
async def _compute_divide_arg():
zip_fn, zip_arg = await asyncio.gather(_compute_zip2_fn(),
_compute_zip2_arg())
return await executor.create_call(zip_fn, zip_arg)
async def _compute_apply_fn():
apply_type = computation_types.FunctionType(
computation_types.StructType(
[divide_blk.type_signature, zip2_type.result]),
computation_types.at_server(divide_blk.type_signature.result))
apply_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_APPLY,
apply_type)
return await executor.create_value(apply_comp, apply_type)
async def _compute_apply_arg():
divide_fn, divide_arg = await asyncio.gather(_compute_divide_fn(),
_compute_divide_arg())
return await executor.create_struct([divide_fn, divide_arg])
async def _compute_divided():
apply_fn, apply_arg = await asyncio.gather(_compute_apply_fn(),
_compute_apply_arg())
return await executor.create_call(apply_fn, apply_arg)
return await _compute_divided()
def _create_compiled_computation(py_fn, parameter_type):
proto, type_signature = tensorflow_computation_factory.create_computation_for_py_fn(
py_fn, parameter_type)
return building_blocks.CompiledComputation(proto,
type_signature=type_signature)
