Add test case for TensorArray concat.
PiperOrigin-RevId: 284087911
Change-Id: I38b82b9d03d85709a89fe8c05ec72ba03023aa49
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 12a2800..ea4f897 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -4811,6 +4811,18 @@
python_version = "PY3",
)
+py_test(
+ name = "tensor_array_ops_test",
+ size = "small",
+ srcs = ["ops/tensor_array_ops_test.py"],
+ python_version = "PY3",
+ deps = [
+ ":array_ops",
+ ":client",
+ ":client_testlib",
+ ],
+)
+
cuda_py_test(
name = "special_math_ops_test",
size = "medium",
diff --git a/tensorflow/python/ops/tensor_array_ops_test.py b/tensorflow/python/ops/tensor_array_ops_test.py
new file mode 100644
index 0000000..4f09ff5
--- /dev/null
+++ b/tensorflow/python/ops/tensor_array_ops_test.py
@@ -0,0 +1,77 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for tensor_array_ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.platform import test
+
+
+class TensorArrayOpsTest(test.TestCase):
+
+ @test_util.run_v1_only('Testing placeholders specifically.')
+ def test_concat_graph(self):
+ values = tensor_array_ops.TensorArray(
+ size=4, dtype=dtypes.string, element_shape=[None], infer_shape=False)
+ a = array_ops.placeholder(dtypes.string, [
+ None,
+ ])
+ b = array_ops.placeholder(dtypes.string, [
+ None,
+ ])
+ values = (values.write(0, a).write(
+ 1, constant_op.constant([], dtypes.string))).write(2, b).write(
+ 3, constant_op.constant([], dtypes.string))
+
+ with self.session() as s:
+ result = s.run(values.concat(), {a: ['a', 'b', 'c'], b: ['c', 'd', 'e']})
+ self.assertAllEqual(result, [b'a', b'b', b'c', b'c', b'd', b'e'])
+
+ @test_util.run_v2_only
+ def test_concat(self):
+ values = tensor_array_ops.TensorArray(
+ size=4, dtype=dtypes.string, element_shape=[None], infer_shape=False)
+ a = constant_op.constant(['a', 'b', 'c'], dtypes.string)
+ b = constant_op.constant(['c', 'd', 'e'], dtypes.string)
+ values = (values.write(0, a).write(
+ 1, constant_op.constant([], dtypes.string))).write(2, b).write(
+ 3, constant_op.constant([], dtypes.string))
+ self.assertAllEqual(values.concat(), [b'a', b'b', b'c', b'c', b'd', b'e'])
+
+ @test_util.run_v2_only
+ def test_concat_in_function(self):
+ @def_function.function
+ def fn(a, b):
+ values = tensor_array_ops.TensorArray(
+ size=4, dtype=dtypes.string, element_shape=[None], infer_shape=False)
+ values = (values.write(0, a).write(
+ 1, constant_op.constant([], dtypes.string))).write(2, b).write(
+ 3, constant_op.constant([], dtypes.string))
+ return values.concat()
+
+ self.assertAllEqual(fn(['a', 'b', 'c'], ['c', 'd', 'e']),
+ [b'a', b'b', b'c', b'c', b'd', b'e'])
+
+
+if __name__ == '__main__':
+ test.main()
