Support constructing TensorFlowVariables from multiple tf operations (ray-project#2182)

Author: decster

python/ray/experimental/tfutils.py

@@ -28,7 +28,7 @@ class TensorFlowVariables(object):
         assignment_nodes (Dict[str, tf.Tensor]): Nodes that assign weights.
     """
 
-    def __init__(self, loss, sess=None, input_variables=None):
+    def __init__(self, output, sess=None, input_variables=None):
         """Creates TensorFlowVariables containing extracted variables.
 
         The variables are extracted by performing a BFS search on the
@@ -38,18 +38,20 @@ def __init__(self, loss, sess=None, input_variables=None):
         variable has a placeholder and assignment operation created for it.
 
         Args:
-            loss (tf.Operation): The tensorflow operation to extract all
-                variables from.
+            output (tf.Operation, List[tf.Operation]): The tensorflow
+                operation to extract all variables from.
             sess (tf.Session): Session used for running the get and set
                 methods.
             input_variables (List[tf.Variables]): Variables to include in the
                 list.
         """
         import tensorflow as tf
         self.sess = sess
-        queue = deque([loss])
+        if not isinstance(output, (list, tuple)):
+            output = [output]
+        queue = deque(output)
         variable_names = []
-        explored_inputs = {loss}
+        explored_inputs = set(output)
 
         # We do a BFS on the dependency graph of the input function to find
         # the variables.

test/tensorflow_test.py

@@ -128,7 +128,7 @@ def testTensorFlowVariables(self):
         variables2.set_flat(flat_weights)
         assert_almost_equal(flat_weights, variables2.get_flat())
 
-        variables3 = ray.experimental.TensorFlowVariables(loss2)
+        variables3 = ray.experimental.TensorFlowVariables([loss2])
         self.assertEqual(variables3.sess, None)
         sess = tf.Session()
         variables3.set_session(sess)