Add support for padding tf.string tensors on CPU.

PiperOrigin-RevId: 188215092

diff --git a/tensorflow/core/kernels/mirror_pad_op.cc b/tensorflow/core/kernels/mirror_pad_op.cc
index 26e1082..1c85c74 100644 (file)
--- a/tensorflow/core/kernels/mirror_pad_op.cc
+++ b/tensorflow/core/kernels/mirror_pad_op.cc
@@ -173,6 +173,7 @@ namespace functor {
   DECLARE_CPU_SPEC(T, int64, 5);
 
 TF_CALL_POD_TYPES(DECLARE_CPU_SPECS);
+TF_CALL_string(DECLARE_CPU_SPECS);
 
 #undef DECLARE_CPU_SPEC
 #undef DECLARE_CPU_SPECS
@@ -194,6 +195,7 @@ TF_CALL_POD_TYPES(DECLARE_CPU_SPECS);
 
 // Note that we do register for bool type, but not in the gradient op.
 TF_CALL_POD_TYPES(REGISTER_KERNEL);
+TF_CALL_string(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 
 #if GOOGLE_CUDA

diff --git a/tensorflow/core/kernels/mirror_pad_op_cpu_impl.h b/tensorflow/core/kernels/mirror_pad_op_cpu_impl.h
index 6716a26..f27ca13 100644 (file)
--- a/tensorflow/core/kernels/mirror_pad_op_cpu_impl.h
+++ b/tensorflow/core/kernels/mirror_pad_op_cpu_impl.h
@@ -29,6 +29,7 @@ using CpuDevice = Eigen::ThreadPoolDevice;
   template struct functor::MirrorPad<CpuDevice, T, int32, CPU_PROVIDED_IXDIM>; \
   template struct functor::MirrorPad<CpuDevice, T, int64, CPU_PROVIDED_IXDIM>;
 TF_CALL_POD_TYPES(DEFINE_CPU_SPECS);
+TF_CALL_string(DEFINE_CPU_SPECS);
 #undef DEFINE_CPU_SPECS
 
 #define DEFINE_CPU_SPECS(T)                                   \

diff --git a/tensorflow/core/kernels/pad_op.cc b/tensorflow/core/kernels/pad_op.cc
index eff3e4d..77c1808 100644 (file)
--- a/tensorflow/core/kernels/pad_op.cc
+++ b/tensorflow/core/kernels/pad_op.cc
@@ -70,7 +70,7 @@ class PadOp : public OpKernel {
             "The first dimension of paddings must be the rank of inputs",
             in1.shape().DebugString(), " ", in0.shape().DebugString()));
 
-    T pad_value(0);
+    T pad_value = T();
     if (context->num_inputs() == 3) {
       const Tensor& constant_values = context->input(2);
       OP_REQUIRES(
@@ -186,6 +186,7 @@ class PadOp : public OpKernel {
                           PadOp<CPUDevice, type, int64>);
 
 TF_CALL_POD_TYPES(REGISTER_KERNEL);
+TF_CALL_string(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 
 #if GOOGLE_CUDA

diff --git a/tensorflow/python/kernel_tests/pad_op_test.py b/tensorflow/python/kernel_tests/pad_op_test.py
index 2c766e3..aaeb3b1 100644 (file)
--- a/tensorflow/python/kernel_tests/pad_op_test.py
+++ b/tensorflow/python/kernel_tests/pad_op_test.py
@@ -238,6 +238,29 @@ class PadOpTest(test.TestCase):
       x = np.random.rand(3, 2, 1, 1).astype(t)
       self._testAll(x + 1j * x, [[0, 0], [0, 0], [0, 0], [0, 0]], 0 + 0j)
 
+  def testString(self):
+    # Numpy does not support padding strings so we compare padding manually.
+    x = ops.convert_to_tensor([["Hello", "World"],
+                               ["Goodnight", "Moon"]])
+
+    constant = array_ops.pad(x, [[1, 0], [0, 1]], mode="CONSTANT",
+                             constant_values="PAD")
+    reflect = array_ops.pad(x, [[1, 0], [0, 1]], mode="REFLECT",
+                            constant_values="PAD")
+    symmetric = array_ops.pad(x, [[1, 0], [0, 1]], mode="SYMMETRIC",
+                              constant_values="PAD")
+    with self.test_session(use_gpu=True):
+      self.assertAllEqual([[b"PAD", b"PAD", b"PAD"],
+                           [b"Hello", b"World", b"PAD"],
+                           [b"Goodnight", b"Moon", b"PAD"]], constant.eval())
+      self.assertAllEqual([[b"Goodnight", b"Moon", b"Goodnight"],
+                           [b"Hello", b"World", b"Hello"],
+                           [b"Goodnight", b"Moon", b"Goodnight"]],
+                          reflect.eval())
+      self.assertAllEqual([[b"Hello", b"World", b"World"],
+                           [b"Hello", b"World", b"World"],
+                           [b"Goodnight", b"Moon", b"Moon"]], symmetric.eval())
+
   def testShapeFunctionEdgeCases(self):
     # Unknown paddings shape.
     inp = constant_op.constant(0.0, shape=[4, 4, 4, 4])