Move WhereOp kernel to the common directory tf2xla/kernels

This will allow it to be used with the new bridge using the fallback path.

Also Added test for the kernel.

PiperOrigin-RevId: 342324585
Change-Id: I343b05809241484ea4fe5081110c7dc773f547ad

tensorflow/compiler/jit/mark_for_compilation_pass.cc

@@ -2040,6 +2040,7 @@ absl::flat_hash_set<string> GetKnownXLAAllowlistOp() {
                                      "UnsortedSegmentSum",
                                      "VarIsInitializedOp",
                                      "VariableShape",
+                                     "Where",
                                      "While",
                                      "XlaBroadcastHelper",
                                      "XlaConv",

tensorflow/compiler/tests/BUILD

@@ -1943,3 +1943,25 @@ tf_xla_py_test(
         "//tensorflow/python:platform_test",
     ],
 )
+
+tf_xla_py_test(
+    name = "where_op_test",
+    size = "small",
+    srcs = ["where_op_test.py"],
+    disabled_backends = [
+        "cpu",
+        "cpu_ondemand",
+        "gpu",
+    ],
+    deps = [
+        ":xla_test",
+        "//tensorflow/compiler/jit",
+        "//tensorflow/contrib/tpu",
+        "//tensorflow/python:array_ops",
+        "//tensorflow/python:client_testlib",
+        "//tensorflow/python:control_flow_ops",
+        "//tensorflow/python:errors",
+        "//tensorflow/python:framework",
+        "//tensorflow/python/compiler/xla:compiler_py",
+    ],
+)

tensorflow/compiler/tests/where_op_test.py

@@ -0,0 +1,90 @@
+# Copyright 2020 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 where op."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+# pylint: enable=g-direct-tensorflow-import
+
+
+class WhereOpTest(xla_test.XLATestCase):
+
+  def testWhere(self):
+    """Test first form of where (return indices)."""
+
+    with self.session() as sess:
+      with self.test_scope():
+        x = array_ops.placeholder(dtypes.bool)
+        true_vals = array_ops.where(x)
+
+      # Output of the computation is dynamic.
+      feed = [[True, False, False], [False, True, True]]
+      self.assertAllEqual([[0, 0], [1, 1], [1, 2]],
+                          sess.run(true_vals, {x: feed}))
+
+  def testWhereGather(self):
+    """Test where followed by a gather."""
+
+    with self.session() as sess:
+      with self.test_scope():
+        x = array_ops.placeholder(dtypes.bool)
+        value = array_ops.constant([[0, 1], [2, 3]], dtypes.float32)
+        true_vals = array_ops.where(x)
+
+        # Gather 0, 2, 3.
+        gathered = array_ops.gather_nd(value, true_vals)
+
+      feed = [[True, False], [True, True]]
+      self.assertAllEqual([0, 2, 3], sess.run(gathered, {x: feed}))
+
+  def testWhereGatherReduce(self):
+    """Test where followed by a gather and a reduce."""
+
+    with self.session() as sess:
+      with self.test_scope():
+        x = array_ops.placeholder(dtypes.bool)
+        value = array_ops.constant([[0, 1], [2, 3]], dtypes.float32)
+        indices = array_ops.where(x)
+
+        # Reduce to 5
+        gathered = array_ops.gather_nd(value, indices)
+        reduction = math_ops.reduce_sum(gathered)
+
+      feed = [[True, False], [True, True]]
+      self.assertAllEqual(5, sess.run(reduction, {x: feed}))
+
+  def testWhere1D(self):
+    """Test first form of where (return indices)."""
+
+    with self.session() as sess:
+      with self.test_scope():
+        x = array_ops.placeholder(dtypes.bool)
+        result = array_ops.where(x)
+
+      # Output of the computation is dynamic.
+      feed = [True, False, True]
+      self.assertAllEqual([[0], [2]], sess.run(result, {x: feed}))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tf2xla/kernels/BUILD

@@ -122,6 +122,7 @@ tf_kernel_library(
         "unary_ops_composition.cc",
         "unpack_op.cc",
         "variable_ops.cc",
+        "where_op.cc",
         "xla_broadcast_helper_op.cc",
         "xla_conv_op.cc",
         "xla_dequantize_op.cc",
@@ -190,6 +191,7 @@ tf_kernel_library(
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/kernels:stateful_random_ops_header",
         "//tensorflow/core/kernels:stateless_random_ops_v2_header",
+        "//tensorflow/core/tpu:tpu_defs",
         "//tensorflow/stream_executor/lib",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",

tensorflow/core/tpu/BUILD

@@ -8,6 +8,7 @@ load(
 
 package(
     default_visibility = [
+        "//tensorflow/compiler/tf2xla/kernels:__subpackages__",
         "//tensorflow/core/tpu:__subpackages__",
         "//tensorflow/stream_executor/tpu:__subpackages__",
     ],

tensorflow/core/tpu/kernels/xla/BUILD

@@ -16,7 +16,6 @@ cc_library(
         "inplace_ops.cc",
         "outfeed_ops.cc",
         "segment_reduction_ops.cc",
-        "where_op.cc",
     ],
     visibility = ["//visibility:public"],
     deps = [