PR comments

Author: phi-dbq

python/sparkdl/graph/utils.py

@@ -62,6 +62,7 @@ def get_op(tfobj_or_name, graph):
                   By default the graph we don't require this argument to be provided.
     """
     graph = validated_graph(graph)
+    _assert_same_graph(tfobj_or_name, graph)
     if isinstance(tfobj_or_name, tf.Operation):
         return tfobj_or_name
     name = tfobj_or_name
@@ -71,8 +72,7 @@ def get_op(tfobj_or_name, graph):
         raise TypeError('invalid op request for [type {}] {}'.format(type(name), name))
     _op_name = op_name(name, graph=None)
     op = graph.get_operation_by_name(_op_name)
-    assert op is not None, \
-        'cannot locate op {} in current graph'.format(_op_name)
+    assert isinstance(op, tf.Operation), 'expect tf.Operation, but got {}'.format(type(op))
     return op
 
 def get_tensor(tfobj_or_name, graph):
@@ -85,6 +85,7 @@ def get_tensor(tfobj_or_name, graph):
                   By default the graph we don't require this argument to be provided.
     """
     graph = validated_graph(graph)
+    _assert_same_graph(tfobj_or_name, graph)
     if isinstance(tfobj_or_name, tf.Tensor):
         return tfobj_or_name
     name = tfobj_or_name
@@ -94,8 +95,7 @@ def get_tensor(tfobj_or_name, graph):
         raise TypeError('invalid tensor request for {} of {}'.format(name, type(name)))
     _tensor_name = tensor_name(name, graph=None)
     tnsr = graph.get_tensor_by_name(_tensor_name)
-    assert tnsr is not None, \
-        'cannot locate tensor {} in current graph'.format(_tensor_name)
+    assert isinstance(tnsr, tf.Tensor), 'expect tf.Tensor, but got {}'.format(type(tnsr))
     return tnsr
 
 def tensor_name(tfobj_or_name, graph=None):
@@ -213,3 +213,10 @@ def strip_and_freeze_until(fetches, graph, sess=None, return_graph=False):
         return g
     else:
         return gdef_frozen
+
+
+def _assert_same_graph(tfobj, graph):
+    if graph is None or not hasattr(tfobj, 'graph'):
+        return
+    assert tfobj.graph == graph, \
+        'the graph of TensorFlow element {} != graph {}'.format(tfobj, graph)

python/tests/graph/test_utils.py

@@ -42,7 +42,31 @@ def _gen_invalid_tensor_or_op_input_with_wrong_types():
         yield TestCase(data=wrong_val, description='wrong type {}'.format(type(wrong_val)))
 
 
-def _gen_valid_tensor_op_objects():
+def _gen_invalid_tensor_or_op_with_graph_pairing():
+    tnsr = tf.constant(1427.08, name='someConstOp')
+    other_graph = tf.Graph()
+    op_name = tnsr.op.name
+
+    # Test get_tensor and get_op returns tensor or op contained in the same graph
+    yield TestCase(data=lambda: tfx.get_op(tnsr, other_graph),
+                   description='test graph from getting op fron tensor')
+    yield TestCase(data=lambda: tfx.get_tensor(tnsr, other_graph),
+                   description='test graph from getting tensor from tensor')
+    yield TestCase(data=lambda: tfx.get_op(tnsr.name, other_graph),
+                   description='test graph from getting op fron tensor name')
+    yield TestCase(data=lambda: tfx.get_tensor(tnsr.name, other_graph),
+                   description='test graph from getting tensor from tensor name')
+    yield TestCase(data=lambda: tfx.get_op(tnsr.op, other_graph),
+                   description='test graph from getting op from op')
+    yield TestCase(data=lambda: tfx.get_tensor(tnsr.op, other_graph),
+                   description='test graph from getting tensor from op')
+    yield TestCase(data=lambda: tfx.get_op(op_name, other_graph),
+                   description='test graph from getting op from op name')
+    yield TestCase(data=lambda: tfx.get_tensor(op_name, other_graph),
+                   description='test graph from getting tensor from op name')
+
+
+def _gen_valid_tensor_op_input_combos():
     op_name = 'someConstOp'
     tnsr_name = '{}:0'.format(op_name)
     tnsr = tf.constant(1427.08, name=op_name)
@@ -154,8 +178,14 @@ def test_invalid_tensor_inputs_with_wrong_types(self, data, description):
         with self.assertRaises(TypeError, msg=description):
             tfx.get_tensor(data, tf.Graph())
 
-    @parameterized.expand(_gen_valid_tensor_op_objects)
+    @parameterized.expand(_gen_valid_tensor_op_input_combos)
     def test_valid_tensor_op_object_inputs(self, data, description):
         """ Must get correct graph elements from valid graph elements or their names """
         tfobj_or_name_a, tfobj_or_name_b = data
         self.assertEqual(tfobj_or_name_a, tfobj_or_name_b, msg=description)
+
+    @parameterized.expand(_gen_invalid_tensor_or_op_with_graph_pairing)
+    def test_invalid_tensor_op_object_graph_pairing(self, data, description):
+        """ Must fail when the graph element is from a different graph than the provided """
+        with self.assertRaises((KeyError, AssertionError, TypeError), msg=description):
+            data()