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_trees_with_one_federated_secure_sum_and_two_federated_aggregates(
self):
federated_aggregate = compiler_test_utils.create_whimsy_called_federated_aggregate(
accumulate_parameter_name='a',
merge_parameter_name='b',
report_parameter_name='c')
federated_secure_sum = compiler_test_utils.create_whimsy_called_federated_secure_sum(
)
called_intrinsics = building_blocks.Struct([
federated_secure_sum,
federated_aggregate,
federated_aggregate,
])
comp = building_blocks.Lambda('d', tf.int32, called_intrinsics)
uri = [
intrinsic_defs.FEDERATED_AGGREGATE.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 test_returns_trees_with_one_federated_secure_sum(self):
federated_secure_sum = compiler_test_utils.create_whimsy_called_federated_secure_sum(
)
called_intrinsics = building_blocks.Struct([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 test_returns_trees_with_one_federated_broadcast(self):
federated_broadcast = compiler_test_utils.create_dummy_called_federated_broadcast(
)
called_intrinsics = building_blocks.Tuple([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 assert_splits_on(self, comp, calls):
"""Asserts that `force_align_and_split_by_intrinsics` removes intrinsics."""
if not isinstance(calls, list):
calls = [calls]
uris = [call.function.uri for call in calls]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, calls)
# Ensure that the resulting computations no longer contain the split
# intrinsics.
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uris))
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uris))
# Removal isn't interesting to test for if it wasn't there to begin with.
self.assertTrue(tree_analysis.contains_called_intrinsic(comp, uris))
if comp.parameter_type is not None:
type_test_utils.assert_types_equivalent(comp.parameter_type,
before.parameter_type)
else:
self.assertIsNone(before.parameter_type)
# THere must be one parameter for each intrinsic in `calls`.
before.type_signature.result.check_struct()
self.assertLen(before.type_signature.result, len(calls))
# Check that `after`'s parameter is a structure like:
# {
# 'original_arg': comp.parameter_type, (if present)
# 'intrinsic_results': [...],
# }
after.parameter_type.check_struct()
if comp.parameter_type is not None:
self.assertLen(after.parameter_type, 2)
type_test_utils.assert_types_equivalent(
comp.parameter_type, after.parameter_type.original_arg)
else:
self.assertLen(after.parameter_type, 1)
# There must be one result for each intrinsic in `calls`.
self.assertLen(after.parameter_type.intrinsic_results, len(calls))
# Check that each pair of (param, result) is a valid type substitution
# for the intrinsic in question.
for i in range(len(calls)):
concrete_signature = computation_types.FunctionType(
before.type_signature.result[i],
after.parameter_type.intrinsic_results[i])
abstract_signature = calls[i].function.intrinsic_def(
).type_signature
type_analysis.check_concrete_instance_of(concrete_signature,
abstract_signature)
def assert_splits_on(self, comp, calls):
"""Asserts that `force_align_and_split_by_intrinsics` removes intrinsics."""
if not isinstance(calls, list):
calls = [calls]
uris = [call.function.uri for call in calls]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, calls)
# Ensure that the resulting computations no longer contain the split
# intrinsics.
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uris))
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uris))
# Removal isn't interesting to test for if it wasn't there to begin with.
self.assertTrue(tree_analysis.contains_called_intrinsic(comp, uris))
self.assert_types_equivalent(comp.parameter_type, before.parameter_type)
# THere must be one parameter for each intrinsic in `calls`.
before.type_signature.result.check_struct()
self.assertLen(before.type_signature.result, len(calls))
# Check that `after`'s parameter is a structure like:
# {
# 'original_arg': comp.parameter_type,
# 'intrinsic_results': [...],
# }
after.parameter_type.check_struct()
self.assertLen(after.parameter_type, 2)
self.assert_types_equivalent(comp.parameter_type,
after.parameter_type.original_arg)
# There must be one result for each intrinsic in `calls`.
self.assertLen(after.parameter_type.intrinsic_results, len(calls))
# Check that each pair of (param, result) is a valid type substitution
# for the intrinsic in question.
for i in range(len(calls)):
concrete_signature = computation_types.FunctionType(
before.type_signature.result[i],
after.parameter_type.intrinsic_results[i])
abstract_signature = calls[i].function.intrinsic_def().type_signature
# `force_align_and_split_by_intrinsics` loses all-equal data due to
# zipping and unzipping. This is okay because the resulting computations
# are not used together directly, but are compiled into unplaced TF code.
abstract_signature = _remove_client_all_equals_from_type(
abstract_signature)
concrete_signature = _remove_client_all_equals_from_type(
concrete_signature)
type_analysis.check_concrete_instance_of(concrete_signature,
abstract_signature)
def test_removes_selected_intrinsic_leaving_remaining_intrinsic(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_secure_sum_bitwidth,
])
comp = building_blocks.Lambda('d', tf.int32, called_intrinsics)
null_aggregate = building_block_factory.create_null_federated_aggregate(
)
secure_sum_bitwidth_uri = federated_secure_sum_bitwidth.function.uri
aggregate_uri = null_aggregate.function.uri
before, after = transformations.force_align_and_split_by_intrinsics(
comp, [null_aggregate])
self.assertTrue(
tree_analysis.contains_called_intrinsic(comp,
secure_sum_bitwidth_uri))
self.assertTrue(
tree_analysis.contains_called_intrinsic(comp, aggregate_uri))
self.assertFalse(
tree_analysis.contains_called_intrinsic(before, aggregate_uri))
self.assertFalse(
tree_analysis.contains_called_intrinsic(after, aggregate_uri))
self.assertTrue(
tree_analysis.contains_called_intrinsic(before,
secure_sum_bitwidth_uri)
or tree_analysis.contains_called_intrinsic(
after, secure_sum_bitwidth_uri))
def test_returns_false_with_unmatched_called_intrinsic(self):
comp = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
uri = intrinsic_defs.FEDERATED_MAP.uri
self.assertFalse(tree_analysis.contains_called_intrinsic(comp, uri))
def test_returns_false_with_no_called_intrinsic(self):
comp = building_block_test_utils.create_identity_function('a')
self.assertFalse(tree_analysis.contains_called_intrinsic(comp))
def test_returns_true_with_matching_uri(self):
comp = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
uri = intrinsic_defs.FEDERATED_BROADCAST.uri
self.assertTrue(tree_analysis.contains_called_intrinsic(comp, uri))
def test_returns_true_with_none_uri(self):
comp = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
self.assertTrue(tree_analysis.contains_called_intrinsic(comp))
def test_raises_type_error_with_none_tree(self):
with self.assertRaises(TypeError):
tree_analysis.contains_called_intrinsic(None)
def get_canonical_form_for_iterative_process(ip):
"""Constructs `tff.backends.mapreduce.CanonicalForm` given iterative process.
This function transforms computations from the input `ip` into
an instance of `tff.backends.mapreduce.CanonicalForm`.
Args:
ip: An instance of `tff.templates.IterativeProcess`.
Returns:
An instance of `tff.backends.mapreduce.CanonicalForm` equivalent to this
process.
Raises:
TypeError: If the arguments are of the wrong types.
transformations.CanonicalFormCompilationError: If the compilation
process fails.
"""
py_typecheck.check_type(ip, iterative_process.IterativeProcess)
initialize_comp = building_blocks.ComputationBuildingBlock.from_proto(
ip.initialize._computation_proto) # pylint: disable=protected-access
next_comp = building_blocks.ComputationBuildingBlock.from_proto(
ip.next._computation_proto) # pylint: disable=protected-access
_check_iterative_process_compatible_with_canonical_form(
initialize_comp, next_comp)
initialize_comp = _replace_intrinsics_with_bodies(initialize_comp)
next_comp = _replace_intrinsics_with_bodies(next_comp)
tree_analysis.check_contains_only_reducible_intrinsics(initialize_comp)
tree_analysis.check_contains_only_reducible_intrinsics(next_comp)
tree_analysis.check_broadcast_not_dependent_on_aggregate(next_comp)
if tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_BROADCAST.uri):
before_broadcast, after_broadcast = (
transformations.force_align_and_split_by_intrinsics(
next_comp, [intrinsic_defs.FEDERATED_BROADCAST.uri]))
else:
before_broadcast, after_broadcast = (
_create_before_and_after_broadcast_for_no_broadcast(next_comp))
contains_federated_aggregate = tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_AGGREGATE.uri)
contains_federated_secure_sum = tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_SECURE_SUM.uri)
if not (contains_federated_aggregate or contains_federated_secure_sum):
raise ValueError(
'Expected an `tff.templates.IterativeProcess` containing at least one '
'`federated_aggregate` or `federated_secure_sum`, found none.')
if contains_federated_aggregate and contains_federated_secure_sum:
before_aggregate, after_aggregate = (
transformations.force_align_and_split_by_intrinsics(
after_broadcast, [
intrinsic_defs.FEDERATED_AGGREGATE.uri,
intrinsic_defs.FEDERATED_SECURE_SUM.uri,
]))
elif contains_federated_secure_sum:
assert not contains_federated_aggregate
before_aggregate, after_aggregate = (
_create_before_and_after_aggregate_for_no_federated_aggregate(
after_broadcast))
else:
assert contains_federated_aggregate and not contains_federated_secure_sum
before_aggregate, after_aggregate = (
_create_before_and_after_aggregate_for_no_federated_secure_sum(
after_broadcast))
type_info = _get_type_info(initialize_comp, before_broadcast, after_broadcast,
before_aggregate, after_aggregate)
initialize = transformations.consolidate_and_extract_local_processing(
initialize_comp)
_check_type_equal(initialize.type_signature, type_info['initialize_type'])
prepare = _extract_prepare(before_broadcast)
_check_type_equal(prepare.type_signature, type_info['prepare_type'])
work = _extract_work(before_aggregate)
_check_type_equal(work.type_signature, type_info['work_type'])
zero, accumulate, merge, report = _extract_federated_aggregate_functions(
before_aggregate)
_check_type_equal(zero.type_signature, type_info['zero_type'])
_check_type_equal(accumulate.type_signature, type_info['accumulate_type'])
_check_type_equal(merge.type_signature, type_info['merge_type'])
_check_type_equal(report.type_signature, type_info['report_type'])
bitwidth = _extract_federated_secure_sum_functions(before_aggregate)
_check_type_equal(bitwidth.type_signature, type_info['bitwidth_type'])
update = _extract_update(after_aggregate)
_check_type_equal(update.type_signature, type_info['update_type'])
return canonical_form.CanonicalForm(
computation_wrapper_instances.building_block_to_computation(initialize),
computation_wrapper_instances.building_block_to_computation(prepare),
computation_wrapper_instances.building_block_to_computation(work),
computation_wrapper_instances.building_block_to_computation(zero),
computation_wrapper_instances.building_block_to_computation(accumulate),
computation_wrapper_instances.building_block_to_computation(merge),
computation_wrapper_instances.building_block_to_computation(report),
computation_wrapper_instances.building_block_to_computation(bitwidth),
computation_wrapper_instances.building_block_to_computation(update))
def get_canonical_form_for_iterative_process(iterative_process):
"""Constructs `tff.backends.mapreduce.CanonicalForm` given iterative process.
This function transforms computations from the input `iterative_process` into
an instance of `tff.backends.mapreduce.CanonicalForm`.
Args:
iterative_process: An instance of `tff.utils.IterativeProcess`.
Returns:
An instance of `tff.backends.mapreduce.CanonicalForm` equivalent to this
process.
Raises:
TypeError: If the arguments are of the wrong types.
transformations.CanonicalFormCompilationError: If the compilation
process fails.
"""
py_typecheck.check_type(iterative_process,
computation_utils.IterativeProcess)
initialize_comp = building_blocks.ComputationBuildingBlock.from_proto(
iterative_process.initialize._computation_proto) # pylint: disable=protected-access
next_comp = building_blocks.ComputationBuildingBlock.from_proto(
iterative_process.next._computation_proto) # pylint: disable=protected-access
_check_iterative_process_compatible_with_canonical_form(
initialize_comp, next_comp)
if len(next_comp.type_signature.result) == 2:
next_comp = _create_next_with_fake_client_output(next_comp)
initialize_comp = _replace_intrinsics_with_bodies(initialize_comp)
next_comp = _replace_intrinsics_with_bodies(next_comp)
tree_analysis.check_intrinsics_whitelisted_for_reduction(initialize_comp)
tree_analysis.check_intrinsics_whitelisted_for_reduction(next_comp)
tree_analysis.check_broadcast_not_dependent_on_aggregate(next_comp)
if tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_BROADCAST.uri):
before_broadcast, after_broadcast = (
transformations.force_align_and_split_by_intrinsics(
next_comp, [intrinsic_defs.FEDERATED_BROADCAST.uri]))
else:
before_broadcast, after_broadcast = (
_create_before_and_after_broadcast_for_no_broadcast(next_comp))
before_aggregate, after_aggregate = (
transformations.force_align_and_split_by_intrinsics(
after_broadcast, [intrinsic_defs.FEDERATED_AGGREGATE.uri]))
type_info = _get_type_info(initialize_comp, before_broadcast,
after_broadcast, before_aggregate,
after_aggregate)
initialize = transformations.consolidate_and_extract_local_processing(
initialize_comp)
_check_type_equal(initialize.type_signature, type_info['initialize_type'])
prepare = _extract_prepare(before_broadcast)
_check_type_equal(prepare.type_signature, type_info['prepare_type'])
work = _extract_work(before_aggregate, after_aggregate)
_check_type_equal(work.type_signature, type_info['work_type'])
zero, accumulate, merge, report = _extract_aggregate_functions(
before_aggregate)
_check_type_equal(zero.type_signature, type_info['zero_type'])
_check_type_equal(accumulate.type_signature, type_info['accumulate_type'])
_check_type_equal(merge.type_signature, type_info['merge_type'])
_check_type_equal(report.type_signature, type_info['report_type'])
update = _extract_update(after_aggregate)
_check_type_equal(update.type_signature, type_info['update_type'])
return canonical_form.CanonicalForm(
computation_wrapper_instances.building_block_to_computation(
initialize),
computation_wrapper_instances.building_block_to_computation(prepare),
computation_wrapper_instances.building_block_to_computation(work),
computation_wrapper_instances.building_block_to_computation(zero),
computation_wrapper_instances.building_block_to_computation(
accumulate),
computation_wrapper_instances.building_block_to_computation(merge),
computation_wrapper_instances.building_block_to_computation(report),
computation_wrapper_instances.building_block_to_computation(update))
def get_canonical_form_for_iterative_process(iterative_process):
"""Constructs `tff.backends.mapreduce.CanonicalForm` given iterative process.
This function transforms computations from the input `iterative_process` into
an instance of `tff.backends.mapreduce.CanonicalForm`.
Args:
iterative_process: An instance of `tff.utils.IterativeProcess`.
Returns:
An instance of `tff.backends.mapreduce.CanonicalForm` equivalent to this
process.
Raises:
TypeError: If the arguments are of the wrong types.
transformations.CanonicalFormCompilationError: If the compilation
process fails.
"""
py_typecheck.check_type(iterative_process,
computation_utils.IterativeProcess)
initialize_comp = building_blocks.ComputationBuildingBlock.from_proto(
iterative_process.initialize._computation_proto) # pylint: disable=protected-access
next_comp = building_blocks.ComputationBuildingBlock.from_proto(
iterative_process.next._computation_proto) # pylint: disable=protected-access
if not (isinstance(next_comp.type_signature.parameter,
computation_types.NamedTupleType)
and isinstance(next_comp.type_signature.result,
computation_types.NamedTupleType)):
raise TypeError(
'Any IterativeProcess compatible with CanonicalForm must '
'have a `next` function which takes and returns instances '
'of `tff.NamedTupleType`; your next function takes '
'parameters of type {} and returns results of type {}'.format(
next_comp.type_signature.parameter,
next_comp.type_signature.result))
if len(next_comp.type_signature.result) == 2:
next_comp = _create_next_with_fake_client_output(next_comp)
initialize_comp = replace_intrinsics_with_bodies(initialize_comp)
next_comp = replace_intrinsics_with_bodies(next_comp)
tree_analysis.check_intrinsics_whitelisted_for_reduction(initialize_comp)
tree_analysis.check_intrinsics_whitelisted_for_reduction(next_comp)
tree_analysis.check_broadcast_not_dependent_on_aggregate(next_comp)
if tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_BROADCAST.uri):
before_broadcast, after_broadcast = (
transformations.force_align_and_split_by_intrinsics(
next_comp, [intrinsic_defs.FEDERATED_BROADCAST.uri]))
else:
before_broadcast, after_broadcast = (
_create_before_and_after_broadcast_for_no_broadcast(next_comp))
before_aggregate, after_aggregate = (
transformations.force_align_and_split_by_intrinsics(
after_broadcast, [intrinsic_defs.FEDERATED_AGGREGATE.uri]))
init_info_packed = pack_initialize_comp_type_signature(
initialize_comp.type_signature)
next_info_packed = pack_next_comp_type_signature(next_comp.type_signature,
init_info_packed)
before_broadcast_info_packed = (
check_and_pack_before_broadcast_type_signature(
before_broadcast.type_signature, next_info_packed))
before_aggregate_info_packed = (
check_and_pack_before_aggregate_type_signature(
before_aggregate.type_signature, before_broadcast_info_packed))
canonical_form_types = check_and_pack_after_aggregate_type_signature(
after_aggregate.type_signature, before_aggregate_info_packed)
initialize = transformations.consolidate_and_extract_local_processing(
initialize_comp)
if initialize.type_signature.result != canonical_form_types[
'initialize_type'].member:
raise transformations.CanonicalFormCompilationError(
'Compilation of initialize has failed. Expected to extract a '
'`building_blocks.CompiledComputation` of type {}, instead we extracted '
'a {} of type {}.'.format(next_comp.type_signature.parameter[0],
type(initialize),
initialize.type_signature.result))
prepare = extract_prepare(before_broadcast, canonical_form_types)
work = extract_work(before_aggregate, after_aggregate,
canonical_form_types)
zero_noarg_function, accumulate, merge, report = extract_aggregate_functions(
before_aggregate, canonical_form_types)
update = extract_update(after_aggregate, canonical_form_types)
cf = canonical_form.CanonicalForm(
computation_wrapper_instances.building_block_to_computation(
initialize),
computation_wrapper_instances.building_block_to_computation(prepare),
computation_wrapper_instances.building_block_to_computation(work),
computation_wrapper_instances.building_block_to_computation(
zero_noarg_function),
computation_wrapper_instances.building_block_to_computation(
accumulate),
computation_wrapper_instances.building_block_to_computation(merge),
computation_wrapper_instances.building_block_to_computation(report),
computation_wrapper_instances.building_block_to_computation(update))
return cf
def get_canonical_form_for_iterative_process(
ip: iterative_process.IterativeProcess,
grappler_config: Optional[
tf.compat.v1.ConfigProto] = _GRAPPLER_DEFAULT_CONFIG):
"""Constructs `tff.backends.mapreduce.CanonicalForm` given iterative process.
This function transforms computations from the input `ip` into
an instance of `tff.backends.mapreduce.CanonicalForm`.
Args:
ip: An instance of `tff.templates.IterativeProcess` that is compatible
with canonical form. Iterative processes are only compatible if: -
`initialize_fn` returns a single federated value placed at `SERVER`. -
`next` takes exactly two arguments. The first must be the state value
placed at `SERVER`. - `next` returns exactly two values.
grappler_config: An optional instance of `tf.compat.v1.ConfigProto` to
configure Grappler graph optimization of the TensorFlow graphs backing the
resulting `tff.backends.mapreduce.CanonicalForm`. These options are
combined with a set of defaults that aggressively configure Grappler. If
`None`, Grappler is bypassed.
Returns:
An instance of `tff.backends.mapreduce.CanonicalForm` equivalent to this
process.
Raises:
TypeError: If the arguments are of the wrong types.
transformations.CanonicalFormCompilationError: If the compilation
process fails.
"""
py_typecheck.check_type(ip, iterative_process.IterativeProcess)
if grappler_config is not None:
py_typecheck.check_type(grappler_config, tf.compat.v1.ConfigProto)
overridden_grappler_config = tf.compat.v1.ConfigProto()
overridden_grappler_config.CopyFrom(_GRAPPLER_DEFAULT_CONFIG)
overridden_grappler_config.MergeFrom(grappler_config)
grappler_config = overridden_grappler_config
initialize_comp = building_blocks.ComputationBuildingBlock.from_proto(
ip.initialize._computation_proto) # pylint: disable=protected-access
next_comp = building_blocks.ComputationBuildingBlock.from_proto(
ip.next._computation_proto) # pylint: disable=protected-access
_check_iterative_process_compatible_with_canonical_form(
initialize_comp, next_comp)
initialize_comp = _replace_intrinsics_with_bodies(initialize_comp)
next_comp = _replace_intrinsics_with_bodies(next_comp)
next_comp = _replace_lambda_body_with_call_dominant_form(next_comp)
tree_analysis.check_contains_only_reducible_intrinsics(initialize_comp)
tree_analysis.check_contains_only_reducible_intrinsics(next_comp)
tree_analysis.check_broadcast_not_dependent_on_aggregate(next_comp)
if tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_BROADCAST.uri):
before_broadcast, after_broadcast = (
transformations.force_align_and_split_by_intrinsics(
next_comp, [intrinsic_defs.FEDERATED_BROADCAST.uri]))
else:
before_broadcast, after_broadcast = (
_create_before_and_after_broadcast_for_no_broadcast(next_comp))
contains_federated_aggregate = tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_AGGREGATE.uri)
contains_federated_secure_sum = tree_analysis.contains_called_intrinsic(
next_comp, intrinsic_defs.FEDERATED_SECURE_SUM.uri)
if not (contains_federated_aggregate or contains_federated_secure_sum):
raise ValueError(
'Expected an `tff.templates.IterativeProcess` containing at least one '
'`federated_aggregate` or `federated_secure_sum`, found none.')
if contains_federated_aggregate and contains_federated_secure_sum:
before_aggregate, after_aggregate = (
transformations.force_align_and_split_by_intrinsics(
after_broadcast, [
intrinsic_defs.FEDERATED_AGGREGATE.uri,
intrinsic_defs.FEDERATED_SECURE_SUM.uri,
]))
elif contains_federated_secure_sum:
assert not contains_federated_aggregate
before_aggregate, after_aggregate = (
_create_before_and_after_aggregate_for_no_federated_aggregate(
after_broadcast))
else:
assert contains_federated_aggregate and not contains_federated_secure_sum
before_aggregate, after_aggregate = (
_create_before_and_after_aggregate_for_no_federated_secure_sum(
after_broadcast))
type_info = _get_type_info(initialize_comp, before_broadcast,
after_broadcast, before_aggregate,
after_aggregate)
initialize = transformations.consolidate_and_extract_local_processing(
initialize_comp, grappler_config)
_check_type_equal(initialize.type_signature, type_info['initialize_type'])
prepare = _extract_prepare(before_broadcast, grappler_config)
_check_type_equal(prepare.type_signature, type_info['prepare_type'])
work = _extract_work(before_aggregate, grappler_config)
_check_type_equal(work.type_signature, type_info['work_type'])
zero, accumulate, merge, report = _extract_federated_aggregate_functions(
before_aggregate, grappler_config)
_check_type_equal(zero.type_signature, type_info['zero_type'])
_check_type_equal(accumulate.type_signature, type_info['accumulate_type'])
_check_type_equal(merge.type_signature, type_info['merge_type'])
_check_type_equal(report.type_signature, type_info['report_type'])
bitwidth = _extract_federated_secure_sum_functions(before_aggregate,
grappler_config)
_check_type_equal(bitwidth.type_signature, type_info['bitwidth_type'])
update = _extract_update(after_aggregate, grappler_config)
_check_type_equal(update.type_signature, type_info['update_type'])
next_parameter_names = (name for (
name, _) in structure.iter_elements(ip.next.type_signature.parameter))
server_state_label, client_data_label = next_parameter_names
return canonical_form.CanonicalForm(
computation_wrapper_instances.building_block_to_computation(
initialize),
computation_wrapper_instances.building_block_to_computation(prepare),
computation_wrapper_instances.building_block_to_computation(work),
computation_wrapper_instances.building_block_to_computation(zero),
computation_wrapper_instances.building_block_to_computation(
accumulate),
computation_wrapper_instances.building_block_to_computation(merge),
computation_wrapper_instances.building_block_to_computation(report),
computation_wrapper_instances.building_block_to_computation(bitwidth),
computation_wrapper_instances.building_block_to_computation(update),
server_state_label=server_state_label,
client_data_label=client_data_label)