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 transform(self, local_computation):
if not self.should_transform(local_computation):
return local_computation, False
generated_tf, _ = generate_tensorflow_for_local_computation(
local_computation)
if not (generated_tf.is_compiled_computation() or
(generated_tf.is_call()
and generated_tf.function.is_compiled_computation())):
raise tree_transformations.TransformationError(
'Failed to generate TensorFlow for a local function. '
f'Generated a building block of type {type(generated_tf)} with '
f'formatted rep {generated_tf.formatted_representation()}.')
return generated_tf, True
def generate_tensorflow_for_local_computation(comp):
"""Generates TensorFlow for a local TFF computation.
This function performs a deduplication of function invocations
according to `tree_analysis.trees_equal`, and hence may reduce the number
of calls under `comp`.
We assume `comp` has type which can be represented by either a call to a
no-arg `building_blocks.CompiledComputation` of type `tensorflow`, or such a
`building_blocks.CompiledComputation` itself. That is, the type signature of
`comp` must be either a potentially nested structure of
`computation_types.TensorType`s and `computation_types.SequenceType`s, or a
function whose parameter and return types are such potentially nested
structures.
Further, we assume that there are no intrinsic or data building blocks inside
`comp`.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` for which we
wish to generate TensorFlow.
Returns:
Either a called instance of `building_blocks.CompiledComputation` or a
`building_blocks.CompiledComputation` itself, depending on whether `comp`
is of non-functional or functional type respectively. Additionally, returns
a boolean to match the `transformation_utils.TransformSpec` pattern.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
names_uniquified, _ = tree_transformations.uniquify_reference_names(comp)
# We ensure the argument to `transform_to_local_call_dominant` is a Lambda, as
# required.
lambda_wrapping_comp = building_blocks.Lambda(None, None, names_uniquified)
# CFG for local CDF plus the type of `lambda_wrapping_comp` imply result must
# be another no-arg lambda.
local_cdf_comp = transform_to_local_call_dominant(
lambda_wrapping_comp).result
def _package_as_deduplicated_block(inner_comp):
repacked_block, _ = tree_transformations.remove_duplicate_block_locals(
inner_comp)
if not repacked_block.is_block():
repacked_block = building_blocks.Block([], repacked_block)
return repacked_block
if local_cdf_comp.type_signature.is_function():
# The CFG for local call dominant tells us that the following patterns are
# possible for a functional computation respecting the structural
# restrictions we require for `comp`:
# 1. CompiledComputation
# 2. Block(bindings, CompiledComp)
# 3. Block(bindings, Lambda(non-functional result with at most one Block))
# 4. Lambda(non-functional result with at most one Block)
if local_cdf_comp.is_compiled_computation():
# Case 1.
return local_cdf_comp, not comp.is_compiled_computation()
elif local_cdf_comp.is_block():
if local_cdf_comp.result.is_compiled_computation():
# Case 2. The bindings in `comp` cannot be referenced in `comp.result`;
# we may return it directly.
return local_cdf_comp.result, True
elif local_cdf_comp.result.is_lambda():
# Case 3. We reduce to case 4 and pass through.
local_cdf_comp = building_blocks.Lambda(
local_cdf_comp.result.parameter_name,
local_cdf_comp.result.parameter_type,
building_blocks.Block(local_cdf_comp.locals,
local_cdf_comp.result.result))
# Reduce potential chain of blocks.
local_cdf_comp, _ = tree_transformations.merge_chained_blocks(
local_cdf_comp)
# This fall-through is intended, since we have merged with case 4.
if local_cdf_comp.is_lambda():
# Case 4.
repacked_block = _package_as_deduplicated_block(
local_cdf_comp.result)
tf_generated, _ = create_tensorflow_representing_block(
repacked_block)
tff_func = building_blocks.Lambda(local_cdf_comp.parameter_name,
local_cdf_comp.parameter_type,
tf_generated)
tf_parser_callable = tree_to_cc_transformations.TFParser()
tf_generated, _ = transformation_utils.transform_postorder(
tff_func, tf_parser_callable)
else:
raise tree_transformations.TransformationError(
'Unexpected structure encountered for functional computation in '
'local call-dominant form: \n'
f'{local_cdf_comp.formatted_representation()}')
else:
# The CFG for local call dominant tells us no lambdas or blocks may be
# present under `comp` for non-functional types which can be represented in
# TensorFlow (in particular, structures of functions are disallowed by this
# restriction). So we may package as a block directly.
repacked_block = _package_as_deduplicated_block(local_cdf_comp)
tf_generated, _ = create_tensorflow_representing_block(repacked_block)
return tf_generated, True