def test_strip_placement_nested_federated_type(self):
int_type = computation_types.TensorType(tf.int32)
server_int_type = computation_types.at_server(int_type)
tupled_int_type = computation_types.to_type((int_type, int_type))
tupled_server_int_type = computation_types.to_type(
(server_int_type, server_int_type))
fed_ref = building_blocks.Reference('x', server_int_type)
before = building_blocks.Struct([fed_ref, fed_ref],
container_type=tuple)
after, modified = tree_transformations.strip_placement(before)
self.assertTrue(modified)
self.assert_has_no_intrinsics_nor_federated_types(after)
type_test_utils.assert_types_identical(before.type_signature,
tupled_server_int_type)
type_test_utils.assert_types_identical(after.type_signature,
tupled_int_type)
def test_returns_trees_with_one_federated_aggregate(self):
federated_aggregate = compiler_test_utils.create_whimsy_called_federated_aggregate(
accumulate_parameter_name='a',
merge_parameter_name='b',
report_parameter_name='c')
called_intrinsics = building_blocks.Struct([federated_aggregate])
comp = building_blocks.Lambda('d', tf.int32, called_intrinsics)
uri = [intrinsic_defs.FEDERATED_AGGREGATE.uri]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, uri)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uri))
self.assertIsInstance(after, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uri))
def test_returns_trees_with_two_federated_broadcast(self):
federated_broadcast = compiler_test_utils.create_dummy_called_federated_broadcast(
)
called_intrinsics = building_blocks.Struct([
federated_broadcast,
federated_broadcast,
])
comp = building_blocks.Lambda('a', tf.int32, called_intrinsics)
uri = [intrinsic_defs.FEDERATED_BROADCAST.uri]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, uri)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uri))
self.assertIsInstance(after, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uri))
def test_returns_trees_with_two_federated_secure_sums(self):
federated_secure_sum = compiler_test_utils.create_whimsy_called_federated_secure_sum(
)
called_intrinsics = building_blocks.Struct([
federated_secure_sum,
federated_secure_sum,
])
comp = building_blocks.Lambda('a', tf.int32, called_intrinsics)
uri = [intrinsic_defs.FEDERATED_SECURE_SUM.uri]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, uri)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uri))
self.assertIsInstance(after, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uri))
def select_output_from_lambda(comp, indices):
"""Constructs a new function with result of selecting `indices` from `comp`.
Args:
comp: Instance of `building_blocks.Lambda` of result type `tff.StructType`
from which we wish to select `indices`. Notice that this named tuple type
must have elements of federated type.
indices: Instance of `str`, `int`, `list`, or `tuple`, specifying the
indices we wish to select from the result of `comp`. If `indices` is a
`str` or `int`, the result of the returned `comp` will be of type at index
`indices` in `comp.type_signature.result`. If `indices` is a `list` or
`tuple`, the result type will be a `tff.StructType` wrapping the specified
selections.
Returns:
A transformed version of `comp` with result value the selection from the
result of `comp` specified by `indices`.
"""
py_typecheck.check_type(comp, building_blocks.Lambda)
py_typecheck.check_type(comp.type_signature.result,
computation_types.StructType)
py_typecheck.check_type(indices, (str, int, tuple, list))
def _select(comp, index):
if comp.is_struct():
return comp[index]
if isinstance(index, str):
return building_blocks.Selection(comp, name=index)
return building_blocks.Selection(comp, index=index)
result_tuple = comp.result
elements = []
if isinstance(indices, (tuple, list)):
for x in indices:
if isinstance(x, (tuple, list)):
selected_output = result_tuple
for y in x:
selected_output = _select(selected_output, y)
else:
selected_output = _select(result_tuple, x)
elements.append(selected_output)
result = building_blocks.Struct(elements)
else:
result = _select(result_tuple, indices)
return building_blocks.Lambda(comp.parameter_name, comp.parameter_type,
result)
def test_value_impl_with_tuple(self):
x_comp = building_blocks.Reference('foo', tf.int32)
y_comp = building_blocks.Reference('bar', tf.bool)
z = value_impl.Value(building_blocks.Struct([x_comp, ('y', y_comp)]))
self.assertIsInstance(z, value_impl.Value)
self.assertEqual(str(z.type_signature), '<int32,y=bool>')
self.assertEqual(str(z), '<foo,y=bar>')
self.assertContainsSubset(['y'], dir(z))
self.assertEqual(str(z.y), 'bar')
self.assertIs(z.y.comp, y_comp)
self.assertLen(z, 2)
self.assertEqual(str(z[0]), 'foo')
self.assertIs(z[0].comp, x_comp)
self.assertEqual(str(z['y']), 'bar')
self.assertIs(z['y'].comp, y_comp)
self.assertEqual(','.join(str(e) for e in iter(z)), 'foo,bar')
with self.assertRaises(SyntaxError):
z(10)
def test_splits_on_multiple_instances_of_two_intrinsics(self):
federated_aggregate = building_block_test_utils.create_whimsy_called_federated_aggregate(
accumulate_parameter_name='a',
merge_parameter_name='b',
report_parameter_name='c')
federated_secure_sum_bitwidth = building_block_test_utils.create_whimsy_called_federated_secure_sum_bitwidth(
)
called_intrinsics = building_blocks.Struct([
federated_aggregate,
federated_aggregate,
federated_secure_sum_bitwidth,
federated_secure_sum_bitwidth,
])
comp = building_blocks.Lambda('d', tf.int32, called_intrinsics)
self.assert_splits_on(comp, [
building_block_factory.create_null_federated_aggregate(),
building_block_factory.create_null_federated_secure_sum_bitwidth()
])
def test_binding_multiple_args_results_in_unique_names(self):
fed_at_clients = computation_types.FederatedType(
tf.int32, placements.CLIENTS)
fed_at_server = computation_types.FederatedType(
tf.int32, placements.SERVER)
tuple_of_federated_types = computation_types.StructType(
[[fed_at_clients], fed_at_server, [fed_at_clients]])
first_selection = building_blocks.Selection(building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=0),
index=0)
second_selection = building_blocks.Selection(building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=2),
index=0)
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Struct([first_selection, second_selection]))
new_lam = form_utils._as_function_of_some_federated_subparameters(
lam, [(0, 0), (2, 0)])
tree_analysis.check_has_unique_names(new_lam)
def test_returns_string_for_struct_with_no_names(self):
data = building_blocks.Data('data', tf.int32)
comp = building_blocks.Struct([data, data])
self.assertEqual(comp.compact_representation(), '<data,data>')
# pyformat: disable
self.assertEqual(
comp.formatted_representation(),
'<\n'
' data,\n'
' data\n'
'>'
)
self.assertEqual(
comp.structural_representation(),
'Struct\n'
'|\n'
'[data, data]'
)
def test_strip_placement_federated_value_at_clients(self):
int_data = building_blocks.Data('x', tf.int32)
float_data = building_blocks.Data('x', tf.float32)
fed_int = building_block_factory.create_federated_value(
int_data, placements.CLIENTS)
fed_float = building_block_factory.create_federated_value(
float_data, placements.CLIENTS)
tup = building_blocks.Struct([fed_int, fed_float],
container_type=tuple)
before = building_block_factory.create_federated_zip(tup)
after, modified = tree_transformations.strip_placement(before)
self.assertTrue(modified)
self.assert_has_no_intrinsics_nor_federated_types(after)
tuple_type = computation_types.StructWithPythonType(
[(None, tf.int32), (None, tf.float32)], tuple)
type_test_utils.assert_types_identical(
before.type_signature, computation_types.at_clients(tuple_type))
type_test_utils.assert_types_identical(after.type_signature,
tuple_type)
def _pack_binary_operator_args(x, y, intrinsics, context_stack):
"""Packs arguments to binary operator into a single arg."""
x_type = x.type_signature
y_type = y.type_signature
if _only_tuple_or_tensor(x_type) and _only_tuple_or_tensor(y_type):
needs_zip = False
else:
if not type_analysis.contains(x_type, lambda t: t.is_federated()):
raise TypeError
needs_zip = True
arg = value_impl.ValueImpl(
building_blocks.Struct([
value_impl.ValueImpl.get_comp(x),
value_impl.ValueImpl.get_comp(y)
]), context_stack)
if needs_zip:
arg = intrinsics.federated_zip(arg)
return arg
def sequence_reduce(self, value, zero, op):
"""Implements `sequence_reduce` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
zero = value_impl.to_value(zero, None, self._context_stack)
op = value_impl.to_value(op, None, self._context_stack)
if value.type_signature.is_sequence():
element_type = value.type_signature.element
else:
py_typecheck.check_type(value.type_signature,
computation_types.FederatedType)
py_typecheck.check_type(value.type_signature.member,
computation_types.SequenceType)
element_type = value.type_signature.member.element
op_type_expected = type_factory.reduction_op(zero.type_signature,
element_type)
if not op_type_expected.is_assignable_from(op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
if value.type_signature.is_sequence():
comp = building_block_factory.create_sequence_reduce(
value, zero, op)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
else:
value_type = computation_types.SequenceType(element_type)
intrinsic_type = computation_types.FunctionType((
value_type,
zero.type_signature,
op.type_signature,
), op.type_signature.result)
intrinsic = building_blocks.Intrinsic(
intrinsic_defs.SEQUENCE_REDUCE.uri, intrinsic_type)
ref = building_blocks.Reference('arg', value_type)
tup = building_blocks.Struct((ref, zero, op))
call = building_blocks.Call(intrinsic, tup)
fn = building_blocks.Lambda(ref.name, ref.type_signature, call)
fn_impl = value_impl.ValueImpl(fn, self._context_stack)
return self.federated_map(fn_impl, value)
def test_replaces_lambda_to_called_tf_block_with_replicated_lambda_arg_with_tf_block_of_same_type(
self):
sum_and_add_one_type = computation_types.StructType(
[tf.int32, tf.int32])
sum_and_add_one = _create_compiled_computation(
lambda x: x[0] + x[1] + 1, sum_and_add_one_type)
int_ref = building_blocks.Reference('x', tf.int32)
tuple_of_ints = building_blocks.Struct((int_ref, int_ref))
summed = building_blocks.Call(sum_and_add_one, tuple_of_ints)
lambda_wrapper = building_blocks.Lambda('x', tf.int32, summed)
parsed, modified = parse_tff_to_tf(lambda_wrapper)
self.assertIsInstance(parsed, building_blocks.CompiledComputation)
self.assertTrue(modified)
# TODO(b/157172423): change to assertEqual when Py container is preserved.
parsed.type_signature.check_equivalent_to(
lambda_wrapper.type_signature)
result = test_utils.run_tensorflow(parsed.proto, 17)
self.assertEqual(35, result)
def test_replaces_lambda_to_called_graph_on_tuple_of_selections_from_arg_with_tf_of_same_type_with_names(
self):
identity_tf_block_type = computation_types.StructType(
[tf.int32, tf.bool])
identity_tf_block = building_block_factory.create_compiled_identity(
identity_tf_block_type)
tuple_ref = building_blocks.Reference('x', [('a', tf.int32),
('b', tf.float32),
('c', tf.bool)])
selected_int = building_blocks.Selection(tuple_ref, index=0)
selected_bool = building_blocks.Selection(tuple_ref, index=2)
created_tuple = building_blocks.Struct([selected_int, selected_bool])
called_tf_block = building_blocks.Call(identity_tf_block,
created_tuple)
lambda_wrapper = building_blocks.Lambda('x', [('a', tf.int32),
('b', tf.float32),
('c', tf.bool)],
called_tf_block)
parsed, modified = parse_tff_to_tf(lambda_wrapper)
exec_lambda = computation_wrapper_instances.building_block_to_computation(
lambda_wrapper)
exec_tf = computation_wrapper_instances.building_block_to_computation(
parsed)
self.assertIsInstance(parsed, building_blocks.CompiledComputation)
self.assertTrue(modified)
self.assertEqual(parsed.type_signature, lambda_wrapper.type_signature)
exec_lambda = computation_wrapper_instances.building_block_to_computation(
lambda_wrapper)
exec_tf = computation_wrapper_instances.building_block_to_computation(
parsed)
self.assertEqual(exec_lambda({
'a': 9,
'b': 10.,
'c': False
}), exec_tf({
'a': 9,
'b': 10.,
'c': False
}))
def _construct_tensorflow_representing_single_local_assignment(
arg_ref, arg_class, previous_output, name_to_output_index):
"""Constructs TensorFlow to represent assignment to a block local in sequence.
Creates a tuple which represents all computations in the block local sequence
depending on those variables which have already been processed, by combining
the elements of `previous_output` with the computations in `arg_class`. Then
generates TensorFlow to capture the logic this tuple encapsulates.
Args:
arg_ref: `building_blocks.Reference` to use in representing
`previous_output` inside the body of the Lambda to be parsed to
TensorFlow. Notice that this is here for name safety.
arg_class: `list` of `building_blocks.ComputationBuildingBlock`s which are
dependent on the block local being processed or any preceding block local;
this should be one of the classes resulting from
`group_block_locals_by_namespace`.
previous_output: The result of parsing previous block local bindings into
functions in the same manner.
name_to_output_index: `dict` mapping block local variables to their index in
the result of the generated TensorFlow. This is used to resolve references
in the computations of `arg_class`, but will not be modified.
Returns:
Called instance of `building_blocks.CompiledComputation` representing
the tuple described above.
"""
pass_through_args = [
building_blocks.Selection(source=arg_ref, index=idx)
for idx, _ in enumerate(previous_output.type_signature)
]
vals_replaced = [
_replace_references_in_comp_with_selections_from_arg(
c, arg_ref, name_to_output_index) for c in arg_class
]
return_tuple = building_blocks.Struct(pass_through_args + vals_replaced)
comp_called = construct_tensorflow_calling_lambda_on_concrete_arg(
arg_ref, return_tuple, previous_output)
return comp_called
def generic_plus(arg):
"""Adds two arguments when possible."""
x = arg[0]
y = arg[1]
_check_top_level_compatibility_with_generic_operators(x, y, 'Generic plus')
if _generic_op_can_be_applied(x, y):
return _apply_generic_op(tf.add, x, y)
# TODO(b/136587334): Push this logic down a level
elif x.type_signature.is_struct():
# This case is needed if federated types are nested deeply.
names = [t[0] for t in anonymous_tuple.iter_elements(x.type_signature)]
added = [
value_impl.ValueImpl.get_comp(generic_plus([x[i], y[i]]))
for i in range(len(names))
]
named_added = building_block_factory.create_named_tuple(
building_blocks.Struct(added), names)
return value_impl.ValueImpl(named_added, context_stack)
else:
raise TypeError('Generic plus encountered unexpected type {}, {}'.format(
x.type_signature, y.type_signature))
def test_ingest_zips_value_when_necessary_to_match_federated_type(self):
# Expects `{<int, int>}@C`
@federated_computation.federated_computation(
computation_types.at_clients((tf.int32, tf.int32)))
def fn(_):
return ()
# This thing will be <{int}@C, {int}@C>
arg = building_blocks.Struct([
building_blocks.Reference(
'x',
computation_types.FederatedType(tf.int32, placements.CLIENTS)),
building_blocks.Reference(
'y',
computation_types.FederatedType(tf.int32, placements.CLIENTS))
])
context = federated_computation_context.FederatedComputationContext(
context_stack_impl.context_stack)
with context_stack_impl.context_stack.install(context):
fn(arg)
def generic_multiply(arg):
"""Multiplies two arguments when possible."""
x = arg[0]
y = arg[1]
_check_top_level_compatibility_with_generic_operators(
x, y, 'Generic multiply')
if _generic_op_can_be_applied(x, y):
return _apply_generic_op(tf.multiply, x, y)
elif x.type_signature.is_struct():
# This case is needed if federated types are nested deeply.
names = [t[0] for t in structure.iter_elements(x.type_signature)]
multiplied = [
value_impl.ValueImpl.get_comp(generic_multiply([x[i], y[i]]))
for i in range(len(names))
]
named_multiplied = building_block_factory.create_named_tuple(
building_blocks.Struct(multiplied), names)
return value_impl.ValueImpl(named_multiplied, context_stack)
else:
raise TypeError(
'Generic multiply encountered unexpected type {}, {}'.format(
x.type_signature, y.type_signature))
def test_broadcast_dependent_on_aggregate_fails_well(self):
mrf = mapreduce_test_utils.get_temperature_sensor_example()
it = form_utils.get_iterative_process_for_map_reduce_form(mrf)
next_comp = it.next.to_building_block()
top_level_param = building_blocks.Reference(next_comp.parameter_name,
next_comp.parameter_type)
first_result = building_blocks.Call(next_comp, top_level_param)
middle_param = building_blocks.Struct([
building_blocks.Selection(first_result, index=0),
building_blocks.Selection(top_level_param, index=1)
])
second_result = building_blocks.Call(next_comp, middle_param)
not_reducible = building_blocks.Lambda(next_comp.parameter_name,
next_comp.parameter_type,
second_result)
not_reducible_it = iterative_process.IterativeProcess(
it.initialize,
computation_wrapper_instances.building_block_to_computation(
not_reducible))
with self.assertRaisesRegex(ValueError, 'broadcast dependent on aggregate'):
form_utils.get_map_reduce_form_for_iterative_process(not_reducible_it)
def test_replaces_lambda_to_unnamed_tuple_of_called_graphs_with_tf_of_same_type(
self):
int_tensor_type = computation_types.TensorType(tf.int32)
int_identity_tf_block = building_block_factory.create_compiled_identity(
int_tensor_type)
float_tensor_type = computation_types.TensorType(tf.float32)
float_identity_tf_block = building_block_factory.create_compiled_identity(
float_tensor_type)
tuple_ref = building_blocks.Reference('x', [tf.int32, tf.float32])
selected_int = building_blocks.Selection(tuple_ref, index=0)
selected_float = building_blocks.Selection(tuple_ref, index=1)
called_int_tf_block = building_blocks.Call(int_identity_tf_block,
selected_int)
called_float_tf_block = building_blocks.Call(float_identity_tf_block,
selected_float)
tuple_of_called_graphs = building_blocks.Struct(
[called_int_tf_block, called_float_tf_block])
lambda_wrapper = building_blocks.Lambda('x', [tf.int32, tf.float32],
tuple_of_called_graphs)
parsed, modified = parse_tff_to_tf(lambda_wrapper)
exec_lambda = computation_wrapper_instances.building_block_to_computation(
lambda_wrapper)
exec_tf = computation_wrapper_instances.building_block_to_computation(
parsed)
self.assertIsInstance(parsed, building_blocks.CompiledComputation)
self.assertTrue(modified)
# TODO(b/157172423): change to assertEqual when Py container is preserved.
parsed.type_signature.check_equivalent_to(
lambda_wrapper.type_signature)
exec_lambda = computation_wrapper_instances.building_block_to_computation(
lambda_wrapper)
exec_tf = computation_wrapper_instances.building_block_to_computation(
parsed)
self.assertEqual(exec_lambda([11, 12.]), exec_tf([11, 12.]))
def test_handles_federated_broadcasts_nested_in_tuple(self):
first_broadcast = compiler_test_utils.create_whimsy_called_federated_broadcast(
)
packed_broadcast = building_blocks.Struct([
building_blocks.Data(
'a',
computation_types.FederatedType(
computation_types.TensorType(tf.int32), placements.SERVER)),
first_broadcast
])
sel = building_blocks.Selection(packed_broadcast, index=0)
second_broadcast = building_block_factory.create_federated_broadcast(sel)
result, _ = compiler_transformations.transform_to_call_dominant(
second_broadcast)
comp = building_blocks.Lambda('a', tf.int32, result)
uri = [intrinsic_defs.FEDERATED_BROADCAST.uri]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, uri)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uri))
self.assertIsInstance(after, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uri))
def test_returns_false_for_tuples_with_different_names(self):
data_1 = building_blocks.Data('data', tf.int32)
tuple_1 = building_blocks.Struct([('a', data_1), ('b', data_1)])
data_2 = building_blocks.Data('data', tf.float32)
tuple_2 = building_blocks.Struct([('c', data_2), ('d', data_2)])
self.assertFalse(tree_analysis.trees_equal(tuple_1, tuple_2))
def test_returns_false_for_tuples_with_different_elements(self):
data_1 = building_blocks.Data('data', tf.int32)
tuple_1 = building_blocks.Struct([data_1, data_1])
data_2 = building_blocks.Data('data', tf.float32)
tuple_2 = building_blocks.Struct([data_2, data_2])
self.assertFalse(tree_analysis.trees_equal(tuple_1, tuple_2))
data_1 = building_blocks.Data('data', tf.int32)
tuple_1 = building_blocks.Struct([data_1, data_1])
data_2 = building_blocks.Data('data', tf.int32)
tuple_2 = building_blocks.Struct([data_2, data_2])
self.assertTrue(tree_analysis.trees_equal(tuple_1, tuple_2))
def test_returns_true_for_identical_graphs_with_nans(self):
tf_comp1 = _create_tensorflow_graph_with_nan()
tf_comp2 = _create_tensorflow_graph_with_nan()
self.assertTrue(tree_analysis.trees_equal(tf_comp1, tf_comp2))
non_aggregation_intrinsics = building_blocks.Struct([
(None,
building_block_test_utils.create_whimsy_called_federated_broadcast()),
(None,
building_block_test_utils.create_whimsy_called_federated_value(
placements.CLIENTS))
])
unit = computation_types.StructType([])
trivial_aggregate = building_block_test_utils.create_whimsy_called_federated_aggregate(
value_type=unit)
trivial_mean = building_block_test_utils.create_whimsy_called_federated_mean(
unit)
trivial_sum = building_block_test_utils.create_whimsy_called_federated_sum(
unit)
# TODO(b/120439632) Enable once federated_mean accepts structured weights.
# trivial_weighted_mean = ...
trivial_secure_sum = building_block_test_utils.create_whimsy_called_federated_secure_sum_bitwidth(
unit)
def to_value(
arg: Any,
type_spec,
context_stack: context_stack_base.ContextStack,
parameter_type_hint=None,
) -> 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, `structure.Struct`s, 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.
parameter_type_hint: An optional `computation_types.Type` or value
convertible to it by `computation_types.to_type` which specifies an
argument type to use in the case that `arg` is a
`function_utils.PolymorphicFunction`.
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)
_check_symbol_binding_context(context_stack.current)
if type_spec is not None:
type_spec = computation_types.to_type(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, function_utils.PolymorphicFunction)):
if isinstance(arg, function_utils.PolymorphicFunction):
if parameter_type_hint is None:
raise TypeError(
'Polymorphic computations cannot be converted to TFF values '
'without a type hint. 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`). If you are a TFF developer and think '
'this should be supported, consider providing `parameter_type_hint` '
'as an argument to the encompassing `to_value` conversion.')
parameter_type_hint = computation_types.to_type(parameter_type_hint)
arg = arg.fn_for_argument_type(parameter_type_hint)
py_typecheck.check_type(arg, computation_base.Computation)
result = ValueImpl(arg.to_compiled_building_block(), context_stack)
elif type_spec is not None and type_spec.is_sequence():
result = _wrap_sequence_as_value(arg, type_spec.element, context_stack)
elif isinstance(arg, structure.Struct):
result = ValueImpl(
building_blocks.Struct([
(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in structure.iter_elements(arg)
]), context_stack)
elif py_typecheck.is_named_tuple(arg):
items = arg._asdict().items() # pytype: disable=attribute-error
result = _dictlike_items_to_value(items, context_stack, type(arg))
elif py_typecheck.is_attrs(arg):
items = attr.asdict(
arg, dict_factory=collections.OrderedDict, recurse=False).items()
result = _dictlike_items_to_value(items, context_stack, type(arg))
elif isinstance(arg, dict):
if isinstance(arg, collections.OrderedDict):
items = arg.items()
else:
items = sorted(arg.items())
result = _dictlike_items_to_value(items, context_stack, type(arg))
elif isinstance(arg, (tuple, list)):
result = ValueImpl(
building_blocks.Struct(
[ValueImpl.get_comp(to_value(x, None, context_stack)) for x in arg],
type(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))
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_spec.is_assignable_from(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
def _dictlike_items_to_value(items, context_stack, container_type) -> ValueImpl:
value = building_blocks.Struct(
[(k, ValueImpl.get_comp(to_value(v, None, context_stack)))
for k, v in items], container_type)
return ValueImpl(value, context_stack)
def test_returns_true_for_tuples(self):
data_1 = building_blocks.Data('data', tf.int32)
tuple_1 = building_blocks.Struct([data_1, data_1])
data_2 = building_blocks.Data('data', tf.int32)
tuple_2 = building_blocks.Struct([data_2, data_2])
self.assertTrue(tree_analysis.trees_equal(tuple_1, tuple_2))
def create_nested_syntax_tree():
r"""Constructs computation with explicit ordering for testing traversals.
The goal of this computation is to exercise each switch
in transform_postorder_with_symbol_bindings, at least all those that recurse.
The computation this function constructs can be represented as below.
Notice that the body of the Lambda *does not depend on the Lambda's
parameter*, so that if we were actually executing this call the argument will
be thrown away.
All leaf nodes are instances of `building_blocks.Data`.
Call
/ \
Lambda('arg') Data('k')
|
Block('y','z')-------------
/ |
['y'=Data('a'),'z'=Data('b')] |
Tuple
/ \
Block('v') Block('x')-------
/ \ | |
['v'=Selection] Data('g') ['x'=Data('h'] |
| |
| |
| Block('w')
| / \
Tuple ------ ['w'=Data('i'] Data('j')
/ \
Block('t') Block('u')
/ \ / \
['t'=Data('c')] Data('d') ['u'=Data('e')] Data('f')
Postorder traversals:
If we are reading Data URIs, results of a postorder traversal should be:
[a, b, c, d, e, f, g, h, i, j, k]
If we are reading locals declarations, results of a postorder traversal should
be:
[t, u, v, w, x, y, z]
And if we are reading both in an interleaved fashion, results of a postorder
traversal should be:
[a, b, c, d, t, e, f, u, g, v, h, i, j, w, x, y, z, k]
Preorder traversals:
If we are reading Data URIs, results of a preorder traversal should be:
[a, b, c, d, e, f, g, h, i, j, k]
If we are reading locals declarations, results of a preorder traversal should
be:
[y, z, v, t, u, x, w]
And if we are reading both in an interleaved fashion, results of a preorder
traversal should be:
[y, z, a, b, v, t, c, d, u, e, f, g, x, h, w, i, j, k]
Since we are also exposing the ability to hook into variable declarations,
it is worthwhile considering the order in which variables are assigned in
this tree. Notice that this order maps neither to preorder nor to postorder
when purely considering the nodes of the tree above. This would be:
[arg, y, z, t, u, v, x, w]
Returns:
An instance of `building_blocks.ComputationBuildingBlock`
satisfying the description above.
"""
data_c = building_blocks.Data('c', tf.float32)
data_d = building_blocks.Data('d', tf.float32)
left_most_leaf = building_blocks.Block([('t', data_c)], data_d)
data_e = building_blocks.Data('e', tf.float32)
data_f = building_blocks.Data('f', tf.float32)
center_leaf = building_blocks.Block([('u', data_e)], data_f)
inner_tuple = building_blocks.Struct([left_most_leaf, center_leaf])
selected = building_blocks.Selection(inner_tuple, index=0)
data_g = building_blocks.Data('g', tf.float32)
middle_block = building_blocks.Block([('v', selected)], data_g)
data_i = building_blocks.Data('i', tf.float32)
data_j = building_blocks.Data('j', tf.float32)
right_most_endpoint = building_blocks.Block([('w', data_i)], data_j)
data_h = building_blocks.Data('h', tf.int32)
right_child = building_blocks.Block([('x', data_h)], right_most_endpoint)
result = building_blocks.Struct([middle_block, right_child])
data_a = building_blocks.Data('a', tf.float32)
data_b = building_blocks.Data('b', tf.float32)
whimsy_outer_block = building_blocks.Block([('y', data_a), ('z', data_b)],
result)
whimsy_lambda = building_blocks.Lambda('arg', tf.float32,
whimsy_outer_block)
whimsy_arg = building_blocks.Data('k', tf.float32)
called_lambda = building_blocks.Call(whimsy_lambda, whimsy_arg)
return called_lambda
def test_ok_on_multiple_no_arg_lambdas(self):
data = building_blocks.Data('x', tf.int32)
lambda_1 = building_blocks.Lambda(None, None, data)
lambda_2 = building_blocks.Lambda(None, None, data)
tup = building_blocks.Struct([lambda_1, lambda_2])
tree_analysis.check_has_unique_names(tup)
def create_tensorflow_representing_block(block):
"""Generates non-duplicated TensorFlow for Block locals binding called graphs.
Assuming that the argument `block` satisfies the following conditions:
1. The local variables in `block` are all called graphs, with arbitrary
arguments.
2. The result of the Block contains tuples, selections and references,
but nothing else.
Then `create_tensorflow_representing_block` will generate a structure, which
may contain tensorflow functions, calls to tensorflow functions, and
references, but which have generated this TensorFlow code without duplicating
work done by referencing the block locals.
Args:
block: Instance of `building_blocks.Block`, whose local variables are all
called instances of `building_blocks.CompiledComputation`, and whose
result contains only instances of `building_blocks.Reference`,
`building_blocks.Selection` or `building_blocks.Struct`.
Returns:
A transformed version of `block`, which has pushed references to the called
graphs in the locals of `block` into TensorFlow.
Raises:
TypeError: If `block` is not an instance of `building_blocks.Block`.
ValueError: If the locals of `block` are anything other than called graphs,
or if the result of `block` contains anything other than selections,
references and tuples.
"""
_check_parameters_for_tf_block_generation(block)
name_generator = building_block_factory.unique_name_generator(block)
def _construct_reference_representing(comp_to_represent):
"""Helper closing over `name_generator` for name safety."""
arg_type = comp_to_represent.type_signature
arg_name = next(name_generator)
return building_blocks.Reference(arg_name, arg_type)
top_level_ref = _get_unbound_ref(block)
named_comp_classes = tree_transformations.group_block_locals_by_namespace(
block)
if top_level_ref:
first_comps = [x[1] for x in named_comp_classes[0]]
tup = building_blocks.Struct([top_level_ref] + first_comps)
graph_tup = _generate_simple_tensorflow(tup)
output_comp = construct_tensorflow_calling_lambda_on_concrete_arg(
top_level_ref, graph_tup, top_level_ref)
name_to_output_index = {top_level_ref.name: 0}
else:
output_comp = building_block_factory.create_compiled_empty_tuple()
name_to_output_index = {}
block_local_names = [x[0] for x in block.locals]
def _update_name_to_output_index(name_class):
"""Helper closing over `name_to_output_index` and `block_local_names`."""
offset = len(name_to_output_index.keys())
for idx, comp_name in enumerate(name_class):
for var_name in block_local_names:
if var_name == comp_name:
name_to_output_index[var_name] = idx + offset
if top_level_ref:
first_names = [x[0] for x in named_comp_classes[0]]
_update_name_to_output_index(first_names)
remaining_comp_classes = named_comp_classes[1:]
else:
remaining_comp_classes = named_comp_classes[:]
for named_comp_class in remaining_comp_classes:
if named_comp_class:
comp_class = [x[1] for x in named_comp_class]
name_class = [x[0] for x in named_comp_class]
arg_ref = _construct_reference_representing(output_comp)
output_comp = _construct_tensorflow_representing_single_local_assignment(
arg_ref, comp_class, output_comp, name_to_output_index)
_update_name_to_output_index(name_class)
arg_ref = _construct_reference_representing(output_comp)
result_replaced = _replace_references_in_comp_with_selections_from_arg(
block.result, arg_ref, name_to_output_index)
comp_called = construct_tensorflow_calling_lambda_on_concrete_arg(
arg_ref, result_replaced, output_comp)
return comp_called, True
