Cleanup some transformation code.
@@ -549,6 +549,16 @@ def merge_tuple_intrinsics(comp):
py_typecheck.check_type(comp,
computation_building_blocks.ComputationBuildingBlock)
+ def _should_transform(comp):
+ uri = (
+ intrinsic_defs.FEDERATED_AGGREGATE.uri,
+ intrinsic_defs.FEDERATED_MAP.uri,
+ )
+ return (isinstance(comp, computation_building_blocks.Tuple) and
+ _is_called_intrinsic(comp[0], uri) and all(
+ _is_called_intrinsic(element, comp[0].function.uri)
+ for element in comp))
+
def _get_comps(comp):
"""Constructs a 2 dimentional Python list of computations.
@@ -670,16 +680,6 @@ def merge_tuple_intrinsics(comp):
args.append(arg)
return args
- def _should_transform(comp):
- uri = (
- intrinsic_defs.FEDERATED_AGGREGATE.uri,
- intrinsic_defs.FEDERATED_MAP.uri,
- )
- return (isinstance(comp, computation_building_blocks.Tuple) and
- _is_called_intrinsic(comp[0], uri) and all(
- _is_called_intrinsic(element, comp[0].function.uri)
- for element in comp))
-
def _transform(comp):
"""Returns a new transformed computation or `comp`."""
if not _should_transform(comp):
@@ -894,10 +894,8 @@ def replace_selection_from_tuple_with_element(comp):
computation_building_blocks.ComputationBuildingBlock)
def _should_transform(comp):
- if (isinstance(comp, computation_building_blocks.Selection) and
- isinstance(comp.source, computation_building_blocks.Tuple)):
- return True
- return False
+ return (isinstance(comp, computation_building_blocks.Selection) and
+ isinstance(comp.source, computation_building_blocks.Tuple))
def _get_index_from_name(selection_name, tuple_type_signature):
type_elements = anonymous_tuple.to_elements(tuple_type_signature)
