Enable unary TensorFlow ops in "xla-legalize-tf-with-tf2xla" pass

This enables unary ops in unary_ops_test.py that have kernel defined in tf2xla/kernels/unary_ops.cc and doesn't already have legalizations. Some tests are disabled if the op is not supported or either using unsigned int or complex constants.

This also deletes unary_mlir_ops_test test now the old and new bridge tests are consolidated.

PiperOrigin-RevId: 306570457
Change-Id: Idbaab1e8986c3659916ae0da28f60bf1960b9f4e

tensorflow/compiler/mlir/xla/transforms/legalize_tf_with_tf2xla.cc

@@ -77,9 +77,13 @@ static bool IsOpWhitelisted(Operation* op) {
   // building valid MLIR using MlirHloBuilder.
   // TODO(hinsu): Drop explicit whitelist when MLIR based bridge is enabled for
   // all tf2xla kernels.
-  return isa<TF::AbsOp>(op) || isa<TF::Atan2Op>(op) || isa<TF::CastOp>(op) ||
+  return isa<TF::AbsOp>(op) || isa<TF::Atan2Op>(op) ||
-         isa<TF::GreaterOp>(op) || isa<TF::InvOp>(op) ||
+         isa<TF::BiasAddGradOp>(op) || isa<TF::CastOp>(op) ||
-         isa<TF::SelectV2Op>(op);
+         isa<TF::ComplexAbsOp>(op) || isa<TF::GreaterOp>(op) ||
+         isa<TF::InvOp>(op) || isa<TF::InvertOp>(op) || isa<TF::LogOp>(op) ||
+         isa<TF::LogicalNotOp>(op) || isa<TF::NegOp>(op) ||
+         isa<TF::SelectV2Op>(op) || isa<TF::SinOp>(op) ||
+         isa<TF::SquareOp>(op) || isa<TF::UnpackOp>(op);
 }
 
 static std::unique_ptr<tensorflow::StaticDeviceMgr> CreateDeviceMgr(

tensorflow/compiler/tests/BUILD

@@ -1346,26 +1346,6 @@ tf_xla_py_test(
     name = "unary_ops_test",
     size = "medium",
     srcs = ["unary_ops_test.py"],
-    python_version = "PY3",
-    tags = [
-        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
-    ],
-    deps = [
-        ":xla_test",
-        "//tensorflow/python:array_ops",
-        "//tensorflow/python:framework",
-        "//tensorflow/python:math_ops",
-        "//tensorflow/python:nn_ops",
-        "//tensorflow/python:nn_ops_gen",
-        "//tensorflow/python:platform_test",
-    ],
-)
-
-# TODO(hinsu): Combine this test with unary_ops_test instead of replicating it.
-tf_xla_py_test(
-    name = "unary_mlir_ops_test",
-    size = "medium",
-    srcs = ["unary_mlir_ops_test.py"],
     enable_mlir_bridge = True,
     python_version = "PY3",
     tags = [

tensorflow/compiler/tests/unary_mlir_ops_test.py

@@ -1,73 +0,0 @@
-# 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 XLA JIT compiler."""
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-
-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 googletest
-
-
-class UnaryOpsTest(xla_test.XLATestCase):
-  """Test cases for unary operators."""
-
-  def _assertOpOutputMatchesExpected(self,
-                                     op,
-                                     inp,
-                                     expected,
-                                     equality_test=None,
-                                     rtol=1e-3,
-                                     atol=1e-5):
-    """Verifies that 'op' produces 'expected' when fed input 'inp' .
-
-    Args:
-      op: operator to test
-      inp: numpy input array to use as input to 'op'.
-      expected: numpy array representing the expected output of 'op'.
-      equality_test: either None, or a function that tests two numpy arrays for
-        equality. If None, self.assertAllClose is used.
-      rtol: relative tolerance for equality test.
-      atol: absolute tolerance for equality test.
-    """
-    with self.session() as session:
-      with self.test_scope():
-        pinp = array_ops.placeholder(
-            dtypes.as_dtype(inp.dtype), inp.shape, name='a')
-        output = op(pinp)
-      result = session.run(output, {pinp: inp})
-      if equality_test is None:
-        self.assertEqual(output.dtype, expected.dtype)
-        self.assertAllCloseAccordingToType(
-            expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03)
-      else:
-        equality_test(result, expected, rtol=rtol, atol=atol)
-
-  def testNumericOps(self):
-    for dtype in self.numeric_types - {np.int8, np.uint8}:
-      self._assertOpOutputMatchesExpected(
-          math_ops.abs,
-          np.array([[2, -1]], dtype=dtype),
-          expected=np.array([[2, 1]], dtype=np.real(dtype(0)).dtype))
-
-
-if __name__ == '__main__':
-  googletest.main()

tensorflow/compiler/tests/unary_ops_test.py

@@ -25,6 +25,7 @@ from six.moves import xrange  # pylint: disable=redefined-builtin
 
 from tensorflow.compiler.tests import xla_test
 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 bitwise_ops
 from tensorflow.python.ops import gen_nn_ops
@@ -84,6 +85,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
     for i in xrange(len(result)):
       self.assertAllClose(result[i], expected[i], rtol, atol)
 
+  @test_util.disable_mlir_bridge(
+      "MlirHloBuilder::Iota missing required for xla::Diag")
   def testAllTypeOps(self):
     for dtype in self.numeric_types - {np.int8, np.uint8}:
       self._assertOpOutputMatchesExpected(
@@ -183,6 +186,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
         self._assertOpOutputMatchesExpected(
             math_ops.cos, x, expected=np.cos(x), rtol=tol, atol=1e-5)
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/153812660): Handle tf.Softmax compilation")
   def testFloatOps(self):
     for dtype in self.float_types:
       x = np.arange(-0.90, 0.90, 0.25)
@@ -593,6 +598,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
           np.array([-1, -0.5, 0, 0.3], dtype=dtype),
           expected=np.array([-1., -0.5, 0., 0.296875], dtype=dtype))
 
+  @test_util.disable_mlir_bridge(
+      "Complex types not supported in CreateDenseElementsAttrFromLiteral")
   def testComplexOps(self):
     for dtype in self.complex_types:
 
@@ -750,6 +757,7 @@ class UnaryOpsTest(xla_test.XLATestCase):
           np.array([1 + 3j, -4 + 7j, 2.7, -3j], dtype=dtype),
           expected=np.array([1, -4, 2.7, 0], dtype=ctypes[dtype]))
 
+  @test_util.disable_mlir_bridge("TODO(b/153896312): Handle unsigned ints")
   def testIntOps(self):
     for dtype in self.int_types:
       self._assertOpOutputMatchesExpected(
@@ -823,6 +831,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
             [[[1., 2.], [3., 4.]], [[5., 6.], [7., 8.]]], dtype=np.float32),
         expected=np.array([14., 22.], dtype=np.float32))
 
+  @test_util.disable_mlir_bridge("TODO(b/153812660): Handle tf.Cast compilation"
+                                )
   def testCast(self):
     shapes = [[], [4], [2, 3], [2, 0, 4]]
     types = {
@@ -870,6 +880,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
             src,
             expected=dst)
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/153812660): Handle tf.Bitcast compilation")
   def testBitcast(self):
     self._assertOpOutputMatchesExpected(
         lambda x: array_ops.bitcast(x, dtypes.int32),
@@ -893,12 +905,16 @@ class UnaryOpsTest(xla_test.XLATestCase):
           np.array([1, 0x100000003f800000], np.int64),
           expected=np.array([1, 0x100000003f800000], np.uint64))
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/153812660): Handle tf.InvertPermutation compilation")
   def testInvertPermutation(self):
     self._assertOpOutputMatchesExpected(
         array_ops.invert_permutation,
         np.array([1, 2, 0], np.int32),
         expected=np.array([2, 0, 1], dtype=np.int32))
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/153812660): Handle tf.InvertPermutation compilation")
   def testInvertPermutationTwiceIsNoop(self):
     self._assertOpOutputMatchesExpected(
         lambda x: array_ops.invert_permutation(array_ops.invert_permutation(x)),
@@ -990,6 +1006,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
         ],
         equality_test=self.ListsAreClose)
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/153812660): Handle tf.DepthToSpace compilation")
   def testDepthToSpace(self):
 
     def make_op(data_format):
@@ -1042,6 +1060,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
                               [[[6, 7], [14, 15]], [[22, 23], [30, 31]]]]],
                             dtype=dtype))
 
+  @test_util.disable_mlir_bridge(
+      "TODO(b/153812660): Handle tf.SpaceToDepth compilation")
   def testSpaceToDepth(self):
 
     def make_op(data_format):
@@ -1101,6 +1121,8 @@ class UnaryOpsTest(xla_test.XLATestCase):
     self._assertOpOutputMatchesExpected(
         nn_ops.softplus, features, expected=expected, rtol=1e-6, atol=9.1e-6)
 
+  @test_util.disable_mlir_bridge(
+      "bf16 type not supported in CreateDenseElementsAttrFromLiteral")
   def testSoftplus(self):
     for dtype in self.float_types:
       self._assertSoftplusMatchesExpected([[-2, 0, 8]], dtype)

tensorflow/python/framework/test_util.py

@@ -1738,18 +1738,15 @@ def enable_tf_xla_constant_folding(description):
   return enable_tf_xla_constant_folding_impl
 
 
-# The description is just for documentation purposes.
+# Updates test function by selectively disabling it.
-def disable_xla(description):
+def _disable_test(execute_func):
 
-  def disable_xla_impl(func):
+  def disable_test_impl(func):
-    """Execute the test method only if xla is not enabled."""
 
     def decorator(func):
 
       def decorated(self, *args, **kwargs):
-        if is_xla_enabled():
+        if execute_func:
-          return
-        else:
           return func(self, *args, **kwargs)
 
       return decorated
@@ -1759,7 +1756,21 @@ def disable_xla(description):
 
     return decorator
 
-  return disable_xla_impl
+  return disable_test_impl
+
+
+# The description is just for documentation purposes.
+def disable_xla(description):  # pylint: disable=unused-argument
+  """Execute the test method only if xla is not enabled."""
+  execute_func = not is_xla_enabled()
+  return _disable_test(execute_func)
+
+
+# The description is just for documentation purposes.
+def disable_mlir_bridge(description):  # pylint: disable=unused-argument
+  """Execute the test method only if MLIR bridge is not enabled."""
+  execute_func = not is_mlir_bridge_enabled()
+  return _disable_test(execute_func)
 
 
 def for_all_test_methods(decorator, *args, **kwargs):
@@ -1791,27 +1802,9 @@ def for_all_test_methods(decorator, *args, **kwargs):
 
 # The description is just for documentation purposes.
 def no_xla_auto_jit(description):  # pylint: disable=unused-argument
-
-  def no_xla_auto_jit_impl(func):
   """This test is not intended to be run with XLA auto jit enabled."""
-
+  execute_func = not is_xla_enabled()
-    def decorator(func):
+  return _disable_test(execute_func)
-
-      def decorated(self, *args, **kwargs):
-        if is_xla_enabled():
-          # Skip test if using XLA is forced.
-          return
-        else:
-          return func(self, *args, **kwargs)
-
-      return decorated
-
-    if func is not None:
-      return decorator(func)
-
-    return decorator
-
-  return no_xla_auto_jit_impl
 
 
 # The description is just for documentation purposes.