diff options
Diffstat (limited to 'tensorflow/python/framework/function_test.py')
| -rw-r--r-- | tensorflow/python/framework/function_test.py | 16 |
1 files changed, 11 insertions, 5 deletions
diff --git a/tensorflow/python/framework/function_test.py b/tensorflow/python/framework/function_test.py index 15e41ba91f..1707f929b8 100644 --- a/tensorflow/python/framework/function_test.py +++ b/tensorflow/python/framework/function_test.py @@ -537,19 +537,25 @@ class FunctionTest(test.TestCase): def testResourceVarAsImplicitInput(self): g = ops.Graph() with g.as_default(), ops.device("cpu:0"): + expected_type = dtypes.float32 + expected_shape = tensor_shape.TensorShape((4, 4)) v = variable_scope.get_variable( - "var", (4, 4), dtypes.float32, use_resource=True) + "var", expected_shape, expected_type, use_resource=True) @function.Defun() def Foo(): - return array_ops.identity(v) + captured = array_ops.identity(v) + self.assertEqual(expected_type, captured.dtype) + self.assertEqual(expected_shape, captured.shape) + return captured, array_ops.shape(captured) - y = v.value() - z = Foo() + expected_val = v.value() + actual_val, actual_shape = Foo() with self.test_session(graph=g): v.initializer.run() - self.assertAllEqual(y.eval(), z.eval()) + self.assertAllEqual(expected_val.eval(), actual_val.eval()) + self.assertAllEqual(expected_shape, actual_shape.eval()) def testDefineErrors(self): with ops.Graph().as_default(): |