diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/list_ops_test.py')
| -rw-r--r-- | tensorflow/python/kernel_tests/list_ops_test.py | 26 |
1 files changed, 26 insertions, 0 deletions
diff --git a/tensorflow/python/kernel_tests/list_ops_test.py b/tensorflow/python/kernel_tests/list_ops_test.py index 0f5607712b..ae413edaec 100644 --- a/tensorflow/python/kernel_tests/list_ops_test.py +++ b/tensorflow/python/kernel_tests/list_ops_test.py @@ -170,6 +170,32 @@ class ListOpsTest(test_util.TensorFlowTestCase): list_ops.tensor_list_pop_back( l_cpu, element_dtype=dtypes.float32)[1]), 2.0) + @test_util.run_in_graph_and_eager_modes + def testCPUGPUCopyNested(self): + if not context.num_gpus(): + return + t = constant_op.constant([1.0, 2.0]) + child_l = list_ops.tensor_list_from_tensor(t, element_shape=scalar_shape()) + l = list_ops.empty_tensor_list( + element_shape=constant_op.constant([], dtype=dtypes.int32), + element_dtype=dtypes.variant) + l = list_ops.tensor_list_push_back(l, child_l) + with context.device("gpu:0"): + l_gpu = array_ops.identity(l) + _, child_l_gpu = list_ops.tensor_list_pop_back( + l_gpu, element_dtype=dtypes.variant) + self.assertAllEqual( + self.evaluate( + list_ops.tensor_list_pop_back( + child_l_gpu, element_dtype=dtypes.float32)[1]), 2.0) + l_cpu = array_ops.identity(l_gpu) + _, child_l_cpu = list_ops.tensor_list_pop_back( + l_cpu, element_dtype=dtypes.variant) + self.assertAllEqual( + self.evaluate( + list_ops.tensor_list_pop_back( + child_l_cpu, element_dtype=dtypes.float32)[1]), 2.0) + def testGraphStack(self): with self.cached_session(): tl = list_ops.empty_tensor_list( |