tff.framework.insert_called_tf_identity_at_leaves

View source on GitHub

Inserts an identity TF graph called on References under comp.

For ease of reasoning about and proving completeness of TFF-to-TF translation capabilities, we will maintain the invariant that we constantly pass up the AST instances of the pattern:

                              Call
                            /      \
        CompiledComputation         Reference

Any block of TFF reducible to TensorFlow must have a functional type signature without nested functions, and therefore we may assume there is a single Reference in the code we are parsing to TF. We continually push logic into the compiled computation as we make our way up the AST, preserving the pattern above; when we hit the lambda that binds this reference, we simply unwrap the call.

To perform this process, we must begin with this pattern; otherwise there may be some arbitrary TFF constructs present between any occurrences of TF and the arguments to which they are applied, e.g. arbitrary selections from and nesting of tuples containing references.

insert_called_tf_identity_at_leaves ensures that the pattern above is present at the leaves of any portion of the TFF AST which is destined to be compiled into TensorFlow; that is, any building_blocks.Reference whose type is compatible with stamping into a TensorFlow graph.

comp Instance of building_blocks.ComputationBuildingBlock whose AST we will traverse, replacing appropriate instances of building_blocks.Reference with graphs representing the identity function of the appropriate type called on the same reference. comp must declare a parameter and result type which are both able to be stamped in to a TensorFlow graph.

A possibly modified version of comp, where any references now have a parent of type building_blocks.Call with function an instance of building_blocks.CompiledComputation.