TF-XLA Bridge ops, tests, and registrations for ReverseOp and ReverseV2Op.

Change: 149488853

Author: tensorflower-gardener

tensorflow/compiler/tests/BUILD

@@ -224,6 +224,17 @@ tf_xla_py_test(
     ],
 )
 
+tf_xla_py_test(
+    name = "reverse_ops_test",
+    size = "small",
+    srcs = ["reverse_ops_test.py"],
+    deps = [
+        ":xla_test",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:framework_for_generated_wrappers",
+    ],
+)
+
 tf_xla_py_test(
     name = "ternary_ops_test",
     size = "small",

tensorflow/compiler/tests/randomized_tests.cc

@@ -1747,6 +1747,30 @@ TEST_F(OpTest, Reshape) {
   });
 }
 
+TEST_F(OpTest, Reverse) {
+  Repeatedly([this]() {
+    std::vector<int64> dims = RandomDims(1);
+    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
+    int64 rank = dims.size();
+    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Reverse")
+                                      .Input(RandomTensor(type, dims))
+                                      .Input(RandomTensor(DT_BOOL, {rank}))
+                                      .Attr("T", DT_FLOAT));
+  });
+}
+
+TEST_F(OpTest, ReverseV2) {
+  Repeatedly([this]() {
+    DataType type = Choose<DataType>({DT_INT32, DT_FLOAT});
+    Tensor data = RandomTensor(type);
+    Tensor indices = RandomReductionIndices(data.dims());
+    ExpectTfAndXlaOutputsAreClose(OpTestBuilder("ReverseV2")
+                                      .Input(data)
+                                      .Input(indices)
+                                      .Attr("T", DT_FLOAT));
+  });
+}
+
 TEST_F(OpTest, Rsqrt) {
   Repeatedly([this]() {
     ExpectTfAndXlaOutputsAreClose(OpTestBuilder("Rsqrt")

tensorflow/compiler/tests/reverse_ops_test.py

@@ -0,0 +1,65 @@
+# Copyright 2017 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.
+# ==============================================================================
+"""Functional tests for XLA Reverse Ops."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import itertools
+import numpy as np
+
+from tensorflow.compiler.tests.xla_test import XLATestCase
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import googletest
+
+
+class ReverseOpsTest(XLATestCase):
+
+  def testReverseOneDim(self):
+    shape = (7, 5, 9, 11)
+    for revdim in range(len(shape)):
+      self._AssertReverseEqual([revdim], shape)
+
+  def testReverseMoreThanOneDim(self):
+    shape = (7, 5, 9, 11)
+    for revdims in itertools.chain.from_iterable(
+        itertools.combinations(range(len(shape)), k)
+        for k in range(2, len(shape)+1)):
+      self._AssertReverseEqual(revdims, shape)
+
+  def _AssertReverseEqual(self, revdims, shape):
+    np.random.seed(120)
+    pval = np.random.randint(0, 100, size=shape).astype(float)
+    with self.test_session():
+      with self.test_scope():
+        p = array_ops.placeholder(dtypes.int32, shape=shape)
+        axis = constant_op.constant(
+            np.array(revdims, dtype=np.int32),
+            shape=(len(revdims),), dtype=dtypes.int32)
+        rval = array_ops.reverse(p, axis).eval({p: pval})
+
+        slices = [
+            slice(-1, None, -1) if d in revdims else slice(None)
+            for d in range(len(shape))]
+      self.assertEqual(
+          pval[slices].flatten().tolist(),
+          rval.flatten().tolist())
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tf2xla/const_analysis.cc

@@ -65,6 +65,8 @@ Status BackwardsConstAnalysis(const Graph& g,
       {"Range", "limit"},
       {"Range", "delta"},
       {"Reshape", "shape"},
+      {"Reverse", "dims"},
+      {"ReverseV2", "axis"},
       {"Slice", "begin"},
       {"Slice", "size"},
       {"Split", "split_dim"},

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -43,6 +43,7 @@ tf_kernel_library(
         "relu_op.cc",
         "reshape_op.cc",
         "retval_op.cc",
+        "reverse_op.cc",
         "select_op.cc",
         "sequence_ops.cc",
         "shape_op.cc",

tensorflow/compiler/tf2xla/kernels/reverse_op.cc

@@ -0,0 +1,110 @@
+/* Copyright 2017 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.
+==============================================================================*/
+
+// XLA-specific reverse Op.
+
+#include "tensorflow/compiler/tf2xla/type_util.h"
+#include "tensorflow/compiler/tf2xla/xla_compilation_device.h"
+#include "tensorflow/compiler/tf2xla/xla_helpers.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/compiler/xla/literal_util.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/register_types.h"
+#include "tensorflow/core/framework/tensor.h"
+
+namespace tensorflow {
+namespace {
+
+class ReverseOp : public XlaOpKernel {
+ public:
+  explicit ReverseOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // r = tf.reverse(x, revdims)
+    const TensorShape x_shape = ctx->InputShape(0);
+    const TensorShape revd_shape = ctx->InputShape(1);
+    // Validate input sizes.
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(revd_shape),
+                errors::InvalidArgument("axes must be a vector, not shape ",
+                                        revd_shape.DebugString()));
+    OP_REQUIRES(ctx, revd_shape.num_elements() == x_shape.dims(),
+                errors::InvalidArgument("axes ", revd_shape.DebugString(),
+                                        " must have same number of elements as"
+                                        " than input tensor has dimensions ",
+                                        x_shape.DebugString(), "."));
+    if (revd_shape.num_elements() == 0) {
+      ctx->SetOutput(0, ctx->Input(0));
+      return;
+    }
+    // ComputationBuilder::Rev() requires concrete values for dimensions arg.
+    xla::Literal lax;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputReshaped(1, {x_shape.dims()}, &lax));
+    std::vector<bool> revdims(x_shape.dims());
+    std::copy(lax.preds().begin(), lax.preds().end(), revdims.begin());
+    std::vector<int64> dimensions;
+
+    for (int d = 0; d < x_shape.dims(); ++d) {
+      if (revdims[d]) {
+        dimensions.push_back(d);
+      }
+    }
+
+    ctx->SetOutput(0, ctx->builder()->Rev(ctx->Input(0), dimensions));
+  }
+};
+
+REGISTER_XLA_OP("Reverse", ReverseOp);
+
+class ReverseV2Op : public XlaOpKernel {
+ public:
+  explicit ReverseV2Op(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    // r = tf.reverse(x, axes)
+    const TensorShape x_shape = ctx->InputShape(0);
+    const TensorShape axes_shape = ctx->InputShape(1);
+    // Validate input sizes.
+    OP_REQUIRES(ctx, TensorShapeUtils::IsVector(axes_shape),
+                errors::InvalidArgument("axes must be a vector, not shape ",
+                                        axes_shape.DebugString()));
+    OP_REQUIRES(ctx, axes_shape.num_elements() <= x_shape.dims(),
+                errors::InvalidArgument("axes ", axes_shape.DebugString(),
+                                        " can not have more elements"
+                                        " than input tensor has dimensions ",
+                                        x_shape.DebugString(), "."));
+    // Reverse is a no-op if axes argument is empty.
+    if (axes_shape.num_elements() == 0) {
+      ctx->SetOutput(0, ctx->Input(0));
+      return;
+    }
+    // ComputationBuilder::Rev() requires concrete values for dimensions arg.
+    std::vector<int64> axes;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInputAsIntVector(1, &axes));
+
+    for (int d = 0; d < axes.size(); ++d) {
+      OP_REQUIRES(ctx, (0 <= axes[d]) && (axes[d] < x_shape.dims()),
+                  errors::InvalidArgument(axes[d], " is out of range [0, ",
+                                          x_shape.dims(), ")."));
+    }
+
+    ctx->SetOutput(0, ctx->builder()->Rev(ctx->Input(0), axes));
+  }
+};
+
+REGISTER_XLA_OP("ReverseV2", ReverseV2Op);
+
+}  // namespace
+}  // namespace tensorflow

tensorflow/compiler/tf2xla/op_registrations.cc

@@ -218,6 +218,10 @@ REGISTER_XLA_KERNEL(DEVICE_CPU_XLA_JIT,
                     Name("Reshape").TypeConstraint("T", kCpuAllTypes));
 REGISTER_XLA_KERNEL(DEVICE_CPU_XLA_JIT, Name("ResourceApplyGradientDescent")
                                             .TypeConstraint("T", kCpuAllTypes));
+REGISTER_XLA_KERNEL(DEVICE_CPU_XLA_JIT,
+                    Name("Reverse").TypeConstraint("T", kCpuAllTypes));
+REGISTER_XLA_KERNEL(DEVICE_CPU_XLA_JIT,
+                    Name("ReverseV2").TypeConstraint("T", kCpuAllTypes));
 REGISTER_XLA_KERNEL(DEVICE_CPU_XLA_JIT,
                     Name("Rsqrt").TypeConstraint("T", kCpuFloatTypes));
 REGISTER_XLA_KERNEL(DEVICE_CPU_XLA_JIT,
@@ -493,6 +497,10 @@ REGISTER_XLA_KERNEL(DEVICE_GPU_XLA_JIT,
                     Name("Reshape").TypeConstraint("T", kGpuAllTypes));
 REGISTER_XLA_KERNEL(DEVICE_GPU_XLA_JIT, Name("ResourceApplyGradientDescent")
                                             .TypeConstraint("T", kGpuAllTypes));
+REGISTER_XLA_KERNEL(DEVICE_GPU_XLA_JIT,
+                    Name("Reverse").TypeConstraint("T", kGpuAllTypes))
+REGISTER_XLA_KERNEL(DEVICE_GPU_XLA_JIT,
+                    Name("ReverseV2").TypeConstraint("T", kGpuAllTypes))
 REGISTER_XLA_KERNEL(DEVICE_GPU_XLA_JIT,
                     Name("Rsqrt").TypeConstraint("T", kGpuFloatTypes));
 REGISTER_XLA_KERNEL(DEVICE_GPU_XLA_JIT,