def test_concatenate_inputs_and_outputs_two_add_one_graphs(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
graph2, input_name_2, output_name_2 = _make_add_one_graph()
with graph1.as_default():
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
with graph2.as_default():
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [input_name_1],
[output_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [input_name_2],
[output_name_2])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, in_name_maps, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0
})
self.assertAllClose(outputs, np.array([2., 3.]))
def test_concatenate_inputs_and_outputs_with_dataset_wires_correctly(self):
dataset_graph, _, dataset_out_name = _make_dataset_constructing_graph()
graph_1, _, out_name_1 = _make_manual_reduce_graph(
dataset_graph, dataset_out_name)
graph_2, _, out_name_2 = _make_manual_reduce_graph(
dataset_graph, dataset_out_name)
with graph_1.as_default():
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
with graph_2.as_default():
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph_1.as_graph_def(),
init_op_name_1, [], [out_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph_2.as_graph_def(),
init_op_name_2, [], [out_name_2])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, _, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
tens = sess.run([
out_name_maps[0][out_name_1], out_name_maps[1][out_name_2]
])
self.assertEqual(tens, [10, 10])
def test_concatenate_inputs_and_outputs_no_arg_graphs(self):
graph1 = tf.Graph()
with graph1.as_default():
out1 = tf.constant(1.0)
init_op_name_1 = tf.compat.v1.global_variables_initializer().name
graph2 = tf.Graph()
with graph2.as_default():
out2 = tf.constant(2.0)
init_op_name_2 = tf.compat.v1.global_variables_initializer().name
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(),
init_op_name_1, [], [out1.name])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(),
init_op_name_2, [], [out2.name])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, _, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with merged_graph.as_default():
with tf.compat.v1.Session() as sess:
sess.run(init_op_name)
outputs = sess.run(
[out_name_maps[0][out1.name], out_name_maps[1][out2.name]])
self.assertAllClose(outputs, np.array([1., 2.]))
def test_concatenate_inputs_and_outputs_no_init_op_graphs(self):
graph1, input_name_1, output_name_1 = _make_add_one_graph()
graph2, input_name_2, output_name_2 = _make_add_one_graph()
graph_spec_1 = graph_spec.GraphSpec(graph1.as_graph_def(), None,
[input_name_1], [output_name_1])
graph_spec_2 = graph_spec.GraphSpec(graph2.as_graph_def(), None,
[input_name_2], [output_name_2])
arg_list = [graph_spec_1, graph_spec_2]
merged_graph, init_op_name, in_name_maps, out_name_maps = graph_merge.concatenate_inputs_and_outputs(
arg_list)
with tf.compat.v1.Session(graph=merged_graph) as sess:
sess.run(init_op_name)
outputs = sess.run(
[
out_name_maps[0][output_name_1],
out_name_maps[1][output_name_2]
],
feed_dict={
in_name_maps[0][input_name_1]: 1.0,
in_name_maps[1][input_name_2]: 2.0
})
self.assertAllClose(outputs, np.array([2., 3.]))
def test_raises_on_non_iterable(self):
with self.assertRaises(TypeError):
graph_merge.concatenate_inputs_and_outputs(1)
def concatenate_tensorflow_blocks(tf_comp_list):
"""Concatenates inputs and outputs of its argument to a single TF block.
Takes a Python `list` or `tuple` of instances of
`computation_building_blocks.CompiledComputation`, and constructs a single
instance of the same building block representing the computations present
in this list concatenated side-by-side.
There is one important convention here for callers to be aware of.
`concatenate_tensorflow_blocks` does not perform any more packing into tuples
than necessary. That is, if `tf_comp_list` contains only a single TF
computation which declares a parameter, the parameter type of the resulting
computation is exactly this single parameter type. Since all TF blocks declare
a result, this is only of concern for parameters, and we will always return a
function with a tuple for its result value.
Args:
tf_comp_list: Python `list` or `tuple` of
`computation_building_blocks.CompiledComputation`s, whose inputs and
outputs we wish to concatenate.
Returns:
A single instance of `computation_building_blocks.CompiledComputation`,
representing all the computations in `tf_comp_list` concatenated
side-by-side.
Raises:
ValueError: If we are passed less than 2 computations in `tf_comp_list`. In
this case, the caller is likely using the wrong function.
TypeError: If `tf_comp_list` is not a `list` or `tuple`, or if it
contains anything other than TF blocks.
"""
py_typecheck.check_type(tf_comp_list, (list, tuple))
if len(tf_comp_list) < 2:
raise ValueError(
'We expect to concatenate at least two blocks of '
'TensorFlow; otherwise the transformation you seek '
'represents simply type manipulation, and you will find '
'your desired function elsewhere in '
'`compiled_computation_transforms`. You passed a tuple of '
'length {}'.format(len(tf_comp_list)))
tf_proto_list = []
for comp in tf_comp_list:
py_typecheck.check_type(
comp, computation_building_blocks.CompiledComputation)
tf_proto_list.append(comp.proto)
(merged_graph, init_op_name, parameter_name_maps,
result_name_maps) = graph_merge.concatenate_inputs_and_outputs(
[_unpack_proto_into_graph_spec(x) for x in tf_proto_list])
concatenated_parameter_bindings = _pack_concatenated_bindings(
[x.tensorflow.parameter for x in tf_proto_list], parameter_name_maps)
concatenated_result_bindings = _pack_concatenated_bindings(
[x.tensorflow.result for x in tf_proto_list], result_name_maps)
if concatenated_parameter_bindings:
tf_result_proto = pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(
merged_graph.as_graph_def()),
initialize_op=init_op_name,
parameter=concatenated_parameter_bindings,
result=concatenated_result_bindings)
else:
tf_result_proto = pb.TensorFlow(
graph_def=serialization_utils.pack_graph_def(
merged_graph.as_graph_def()),
initialize_op=init_op_name,
result=concatenated_result_bindings)
parameter_type = _construct_concatenated_type(
[x.type_signature.parameter for x in tf_comp_list])
return_type = _construct_concatenated_type(
[x.type_signature.result for x in tf_comp_list])
function_type = computation_types.FunctionType(parameter_type, return_type)
serialized_function_type = type_serialization.serialize_type(function_type)
constructed_proto = pb.Computation(type=serialized_function_type,
tensorflow=tf_result_proto)
return computation_building_blocks.CompiledComputation(constructed_proto)
