def test_binds_single_argument_to_lower_lambda(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.NamedTupleType(
[fed_at_clients, fed_at_server])
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Selection(building_blocks.Reference(
'x', tuple_of_federated_types),
index=0))
zeroth_index_extracted = mapreduce_transformations.bind_single_selection_as_argument_to_lower_level_lambda(
lam, 0)
self.assertEqual(zeroth_index_extracted.type_signature,
lam.type_signature)
self.assertIsInstance(zeroth_index_extracted, building_blocks.Lambda)
self.assertIsInstance(zeroth_index_extracted.result,
building_blocks.Call)
self.assertIsInstance(zeroth_index_extracted.result.function,
building_blocks.Lambda)
self.assertRegex(str(zeroth_index_extracted.result.function),
r'\((.{4})1 -> (\1)1\)')
self.assertEqual(str(zeroth_index_extracted.result.argument),
'_var1[0]')
def _extract_prepare(before_broadcast):
"""extracts `prepare` from `before_broadcast`.
This function is intended to be used by
`get_canonical_form_for_iterative_process` only. As a result, this function
does not assert that `before_broadcast` has the expected structure, the
caller is expected to perform these checks before calling this function.
Args:
before_broadcast: The first result of splitting `next_comp` on
`intrinsic_defs.FEDERATED_BROADCAST`.
Returns:
`prepare` as specified by `canonical_form.CanonicalForm`, an instance of
`building_blocks.CompiledComputation`.
Raises:
transformations.CanonicalFormCompilationError: If we extract an AST of the
wrong type.
"""
s1_index_in_before_broadcast = 0
s1_to_s2_computation = (
transformations.bind_single_selection_as_argument_to_lower_level_lambda(
before_broadcast, s1_index_in_before_broadcast)).result.function
return transformations.consolidate_and_extract_local_processing(
s1_to_s2_computation)
def _extract_prepare(before_broadcast):
"""Converts `before_broadcast` into `prepare`.
Args:
before_broadcast: The first result of splitting `next_comp` on
`intrinsic_defs.FEDERATED_BROADCAST`.
Returns:
`prepare` as specified by `canonical_form.CanonicalForm`, an instance of
`building_blocks.CompiledComputation`.
Raises:
transformations.CanonicalFormCompilationError: If we extract an AST of the
wrong type.
"""
# See `get_iterative_process_for_canonical_form()` above for the meaning of
# variable names used in the code below.
s1_index_in_before_broadcast = 0
s1_to_s2_computation = (
transformations.bind_single_selection_as_argument_to_lower_level_lambda(
before_broadcast, s1_index_in_before_broadcast)).result.function
prepare = transformations.consolidate_and_extract_local_processing(
s1_to_s2_computation)
return prepare
def extract_prepare(before_broadcast, canonical_form_types):
"""Converts `before_broadcast` into `prepare`.
Args:
before_broadcast: The first result of splitting `next_comp` on
`intrinsic_defs.FEDERATED_BROADCAST`.
canonical_form_types: `dict` holding the `canonical_form.CanonicalForm` type
signatures specified by the `tff.utils.IterativeProcess` we are compiling.
Returns:
`prepare` as specified by `canonical_form.CanonicalForm`, an instance of
`building_blocks.CompiledComputation`.
Raises:
transformations.CanonicalFormCompilationError: If we fail to extract a
`building_blocks.CompiledComputation`, or we extract one of the wrong type.
"""
# See `get_iterative_process_for_canonical_form()` above for the meaning of
# variable names used in the code below.
s1_index_in_before_broadcast = 0
s1_to_s2_computation = (
transformations.bind_single_selection_as_argument_to_lower_level_lambda(
before_broadcast, s1_index_in_before_broadcast)).result.function
prepare = transformations.consolidate_and_extract_local_processing(
s1_to_s2_computation)
if not isinstance(prepare, building_blocks.CompiledComputation):
raise transformations.CanonicalFormCompilationError(
'Failed to extract a `building_blocks.CompiledComputation` from '
'prepare, instead received a {} (of type {}).'.format(
type(prepare), prepare.type_signature))
if prepare.type_signature != canonical_form_types['prepare_type']:
raise transformations.CanonicalFormCompilationError(
'Extracted a TF block of the wrong type. Expected a function with type '
'{}, but the type signature of the TF block was {}'.format(
canonical_form_types['prepare_type'], prepare.type_signature))
return prepare
