def remove_called_lambdas_and_blocks(comp):
"""Removes any called lambdas and blocks from `comp`.
This function first resolves any higher-order functions, so that replacing
called lambdas with blocks and then inlining the block locals cannot result
in more called lambdas. It then performs this sequence of transformations,
taking care to inline selections from tuples at appropriate stages to prevent
possible combinatorial growth of the generated AST.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` from which we
want to remove called lambdas and blocks.
Returns:
A transformed version of `comp` which has no called lambdas or blocks, and
no extraneous selections from tuples.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
comp, names_uniquified = tree_transformations.uniquify_reference_names(
comp)
# TODO(b/162888191): Remove this gating when `resolve_higher_order_functions`
# is more efficient, or avoided.
if _contains_higher_order_fns(comp):
# `resolve_higher_order_functions` can be expensive, so we only call into it
# when necessary.
comp, fns_resolved = tree_transformations.resolve_higher_order_functions(
comp)
else:
# We must still inline any functional references. We first inline selections
# from tuples to prevent the AST from becoming unnecessarly large.
comp, sels_removed = tree_transformations.inline_selections_from_tuple(
comp)
if sels_removed:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, locals_inlined = tree_transformations.inline_block_locals(comp)
fns_resolved = sels_removed or locals_inlined
comp, lambdas_replaced = tree_transformations.replace_called_lambda_with_block(
comp)
if fns_resolved or lambdas_replaced:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, sels_removed = tree_transformations.inline_selections_from_tuple(
comp)
if sels_removed:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, locals_inlined = tree_transformations.inline_block_locals(comp)
if locals_inlined:
# Inlining local symbols may reintroduce selection-from-tuple pattern,
# combinatorially increasing build times in the worst case. We ensure
# here that remove_called_lambdas_and_blocks respects the postcondition that
# selections from tuples are always collapsed.
comp, _ = tree_transformations.inline_selections_from_tuple(comp)
comp, _ = tree_transformations.uniquify_reference_names(comp)
modified = names_uniquified or fns_resolved or lambdas_replaced or sels_removed or locals_inlined
return comp, modified
def remove_called_lambdas_and_blocks(comp):
"""Removes any called lambdas and blocks from `comp`.
This function first resolves any higher-order functions, so that replacing
called lambdas with blocks and then inlining the block locals cannot result
in more called lambdas. It then performs this sequence of transformations,
taking care to inline selections from tuples before inlining the rest of
the block locals to prevent possible combinatorial growth of the generated
AST.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` from which we
want to remove called lambdas and blocks.
Returns:
A transformed version of `comp` which has no called lambdas or blocks, and
no extraneous selections from tuples.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
comp, names_uniquified = tree_transformations.uniquify_reference_names(comp)
comp, fns_resolved = tree_transformations.resolve_higher_order_functions(comp)
comp, lambdas_replaced = tree_transformations.replace_called_lambda_with_block(
comp)
if fns_resolved or lambdas_replaced:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, sels_removed = tree_transformations.inline_selections_from_tuple(comp)
if sels_removed:
comp, _ = tree_transformations.uniquify_reference_names(comp)
comp, locals_inlined = tree_transformations.inline_block_locals(comp)
modified = names_uniquified or fns_resolved or lambdas_replaced or sels_removed or locals_inlined
return comp, modified
def parse_tff_to_tf(comp):
comp, _ = tree_transformations.insert_called_tf_identity_at_leaves(comp)
parser_callable = tree_to_cc_transformations.TFParser()
comp, _ = tree_transformations.replace_called_lambda_with_block(comp)
comp, _ = tree_transformations.inline_block_locals(comp)
comp, _ = tree_transformations.replace_selection_from_tuple_with_element(comp)
new_comp, transformed = transformation_utils.transform_postorder(
comp, parser_callable)
return new_comp, transformed
def _inline_block_variables_required_to_align_intrinsics(comp, uri):
"""Inlines the variables required to align the intrinsic for the given `uri`.
This function inlines only the block variables required to align an intrinsic,
which is necessary because many transformations insert block variables that do
not impact alignment and should not be inlined.
Additionally, this function iteratively attempts to inline block variables a
long as the intrinsic can not be extracted to the top level lambda. Meaning,
that unbound references in variables that are inlined, will also be inlined.
Args:
comp: The `building_blocks.Lambda` to transform.
uri: A Python `list` of URI of intrinsics.
Returns:
A new computation with the transformation applied or the original `comp`.
Raises:
ValueError: If an there are unbound references, other than block variables,
preventing an intrinsic with the given `uri` from being aligned.
"""
py_typecheck.check_type(comp, building_blocks.Lambda)
py_typecheck.check_type(uri, list)
for x in uri:
py_typecheck.check_type(x, str)
while not _can_extract_intrinsics_to_top_level_lambda(comp, uri):
unbound_references = transformation_utils.get_map_of_unbound_references(
comp)
variable_names = set()
intrinsics = _get_called_intrinsics(comp, uri)
for intrinsic in intrinsics:
names = unbound_references[intrinsic]
names.discard(comp.parameter_name)
variable_names.update(names)
if not variable_names:
raise tree_transformations.TransformationError(
'Inlining `Block` variables has failed. Expected to find unbound '
'references for called `Intrisic`s matching the URI: \'{}\', but '
'none were found in the AST: \n{}'.format(
uri, comp.formatted_representation()))
comp, modified = tree_transformations.inline_block_locals(
comp, variable_names=variable_names)
if modified:
comp, _ = tree_transformations.uniquify_reference_names(comp)
else:
raise tree_transformations.TransformationError(
'Inlining `Block` variables has failed, this will result in an '
'infinite loop. Expected to modify the AST by inlining the variable '
'names: \'{}\', but no transformations to the AST: \n{}'.
format(variable_names, comp.formatted_representation()))
return comp
def _inline_functions(comp):
function_type_reference_names = []
def _populate_function_type_ref_names(comp):
if comp.is_reference() and comp.type_signature.is_function():
function_type_reference_names.append(comp.name)
return comp, False
transformation_utils.transform_postorder(comp,
_populate_function_type_ref_names)
return tree_transformations.inline_block_locals(
comp, variable_names=set(function_type_reference_names))
def transformation_fn(x): x, _ = tree_transformations.remove_mapped_or_applied_identity(x) x, _ = tree_transformations.inline_block_locals(x) x, _ = tree_transformations.replace_selection_from_tuple_with_element(x) return x
def transformation_fn(x):
x, _ = tree_transformations.uniquify_reference_names(x)
x, _ = tree_transformations.inline_block_locals(x)
x, _ = tree_transformations.remove_mapped_or_applied_identity(x)
return x
