Add basic tensorflow subgraph (function) support

[scottcjt]

Hi,

This patch adds basic FunctionDef support in TF. Basically a FunctionDef is a GraphDef with some minor differences such as:

• graph inputs/outputs are explicitly declared
• when referencing a node output, the format is <node_name>:<output_arg_name>:<port>. The output_arg_name is determined by op type, and it is hidden in this patch (which is not a good thing IMHO...).

[lutzroeder]

@scottcjt Can you provide a sample file to test this change?

[scottcjt]

Sure. But I don't know any public url-reachable model using this, so here is a sample graph. tf_function.tar.gz

The most easy way to get a TF graph with functions is to use TF-2.0 and its new control flow. Sample above is generated by:

import tensorflow as tf

def save_to_pb_and_txt(name, gdef):
    from google.protobuf.text_format import MessageToString

    with open('%s.pb' % name, mode='wb') as fp:
        fp.write(gdef.SerializeToString())

    with open('%s.pbtxt' % name, mode='w') as fp:
        fp.write(MessageToString(gdef))

def build_simple_loop_graph():
    @tf.function(input_signature=[tf.TensorSpec(shape=[1], dtype=tf.int32)])
    def main(n):
        total = 0
        i = 0
        while i < n[0]:
            total += i
            i += 1

        return tf.identity(total, name='out')

    # Test...
    result = main([10])
    assert result.numpy() == sum(range(10)), 'expect %d, got %d' % (sum(range(10)), result)

    cfunc = main.get_concrete_function()
    save_to_pb_and_txt('xd', cfunc.graph.as_graph_def())

build_simple_loop_graph()

[lutzroeder]

#342 #168 #68 #320