chore: import upstream snapshot with attribution

tensorflow/compiler/tests/adadelta_test.py

@@ -0,0 +1,138 @@
+# Copyright 2018 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 Adadelta Optimizer."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adadelta
+
+
+class AdadeltaOptimizerTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    num_updates = 4  # number of ADADELTA steps to perform
+    if "CPU" in self.device:
+      # To avoid timeout on CPU.
+      all_grad = [0.2, 0.01]
+      all_lr = [1.0, 0.1]
+    else:
+      all_grad = [0.2, 0.1, 0.01]
+      all_lr = [1.0, 0.5, 0.1]
+
+    for dtype in self.float_types | self.complex_types:
+      with self.session(), self.test_scope():
+        for grad in all_grad:
+          for lr in all_lr:
+            var0_init = [1.0, 2.0]
+            var1_init = [3.0, 4.0]
+            var0 = resource_variable_ops.ResourceVariable(
+                var0_init, dtype=dtype)
+            var1 = resource_variable_ops.ResourceVariable(
+                var1_init, dtype=dtype)
+
+            grads = constant_op.constant([grad, grad], dtype=dtype)
+
+            accum = 0.0
+            accum_update = 0.0
+
+            # ADADELTA gradient optimizer
+            rho = 0.95
+            epsilon = 1e-8
+            adadelta_opt = adadelta.AdadeltaOptimizer(
+                learning_rate=lr, rho=rho, epsilon=epsilon)
+            adadelta_update = adadelta_opt.apply_gradients(
+                zip([grads, grads], [var0, var1]))
+            self.evaluate(variables.global_variables_initializer())
+            opt_vars = adadelta_opt.variables()
+            self.assertStartsWith(opt_vars[0].name, var0._shared_name)
+            self.assertStartsWith(opt_vars[1].name, var0._shared_name)
+            self.assertStartsWith(opt_vars[2].name, var1._shared_name)
+            self.assertStartsWith(opt_vars[3].name, var1._shared_name)
+            self.assertEqual(4, len(opt_vars))
+            # Assign slots
+            slot = [None] * 2
+            slot_update = [None] * 2
+            self.assertEqual(["accum", "accum_update"],
+                             adadelta_opt.get_slot_names())
+            slot[0] = adadelta_opt.get_slot(var0, "accum")
+            self.assertEqual(slot[0].get_shape(), var0.get_shape())
+            self.assertNotIn(slot[0], variables.trainable_variables())
+
+            slot_update[0] = adadelta_opt.get_slot(var0, "accum_update")
+            self.assertEqual(slot_update[0].get_shape(), var0.get_shape())
+            self.assertNotIn(slot_update[0], variables.trainable_variables())
+
+            slot[1] = adadelta_opt.get_slot(var1, "accum")
+            self.assertEqual(slot[1].get_shape(), var1.get_shape())
+            self.assertNotIn(slot[1], variables.trainable_variables())
+
+            slot_update[1] = adadelta_opt.get_slot(var1, "accum_update")
+            self.assertEqual(slot_update[1].get_shape(), var1.get_shape())
+            self.assertNotIn(slot_update[1], variables.trainable_variables())
+
+            # Fetch params to validate initial values
+            self.assertAllClose(var0_init, self.evaluate(var0))
+            self.assertAllClose(var1_init, self.evaluate(var1))
+
+          update = [None] * num_updates
+          tot_update = 0
+          for step in range(num_updates):
+            # Run adadelta update for comparison
+            self.evaluate(adadelta_update)
+
+            # Perform initial update without previous accum values
+            accum = accum * rho + (grad**2) * (1 - rho)
+            update[step] = (
+                np.sqrt(accum_update + epsilon) *
+                (1. / np.sqrt(accum + epsilon)) * grad)
+            accum_update = (
+                accum_update * rho + (update[step]**2) * (1.0 - rho))
+            tot_update += update[step] * lr
+
+            # Check that the accumulators have been updated
+            for slot_idx in range(2):
+              self.assertAllCloseAccordingToType(
+                  np.array([accum, accum], dtype=dtype),
+                  self.evaluate(slot[slot_idx]),
+                  rtol=1e-5)
+
+              self.assertAllCloseAccordingToType(
+                  np.array([accum_update, accum_update], dtype=dtype),
+                  self.evaluate(slot_update[slot_idx]),
+                  rtol=1e-5)
+
+            # Check that the parameters have been updated
+            self.assertAllCloseAccordingToType(
+                np.array(
+                    [var0_init[0] - tot_update, var0_init[1] - tot_update],
+                    dtype=dtype),
+                self.evaluate(var0),
+                rtol=1e-5)
+
+            self.assertAllCloseAccordingToType(
+                np.array(
+                    [var1_init[0] - tot_update, var1_init[1] - tot_update],
+                    dtype=dtype),
+                self.evaluate(var1),
+                rtol=1e-5)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/adagrad_da_test.py

@@ -0,0 +1,161 @@
+# Copyright 2016 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 AdagradDA optimizer."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad_da
+
+
+class AdagradDAOptimizerTest(xla_test.XLATestCase):
+
+  def testAdagradDAWithoutRegularizationBasic1(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        global_step = resource_variable_ops.ResourceVariable(
+            0, dtype=dtypes.int64)
+        var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.0,
+            l2_regularization_strength=0.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+        self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+        # Run a step of AdagradDA
+        update.run()
+
+        # Let g be the gradient accumulator, gg be the gradient squared
+        # accumulator, T be the global step, lr be the learning rate,
+        # and k the initial gradient squared accumulator value.
+        # w = \dfrac{sign(-g)*lr*|g - l1*T|_{+}}{l2*T*lr + \sqrt{k+gg})}
+        # For -0.1*3.0*(0.1 - 0)/(0 + sqrt(0.1 + 0.1*0.1)) = -0.904534
+        # similarly for others.
+        self.assertAllCloseAccordingToType(
+            np.array([-0.904534, -1.603567]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([-0.094821, -0.189358]), self.evaluate(var1))
+
+  def testAdagradDAwithoutRegularizationBasic2(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        global_step = resource_variable_ops.ResourceVariable(
+            0, dtype=dtypes.int64)
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.0,
+            l2_regularization_strength=0.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1))
+
+        # Run a step of AdagradDA
+        update.run()
+
+        self.assertAllCloseAccordingToType(
+            np.array([-0.904534, -1.603567]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([-0.094821, -0.189358]), self.evaluate(var1))
+
+  def testAdagradDAWithL1(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        global_step = resource_variable_ops.ResourceVariable(
+            0, dtype=dtypes.int64)
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=0.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1))
+
+        # Run a step of AdagradDA
+        update.run()
+
+        self.assertAllCloseAccordingToType(
+            np.array([-0.895489, -1.59555]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([-0.085339, -0.17989]), self.evaluate(var1))
+
+  def testAdagradDAWithL1_L2(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        global_step = resource_variable_ops.ResourceVariable(
+            0, dtype=dtypes.int64)
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+
+        opt = adagrad_da.AdagradDAOptimizer(
+            3.0,
+            global_step,
+            initial_gradient_squared_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=2.0)
+        update = opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]), global_step=global_step)
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1))
+
+        # Run a step of AdagradDA
+        update.run()
+
+        self.assertAllCloseAccordingToType(
+            np.array([-0.046907, -0.093659]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([-0.004275, -0.009023]), self.evaluate(var1))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/adagrad_test.py

@@ -0,0 +1,124 @@
+# 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.
+# ==============================================================================
+"""Tests for Adagrad."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+
+
+class AdagradOptimizerTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    for dtype in self.float_types | self.complex_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        ada_opt = adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1)
+        ada_update = ada_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Run 3 steps of adagrad
+        for _ in range(3):
+          ada_update.run()
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-1.6026098728179932, -0.6026098728179932]),
+            self.evaluate(var0),
+            float_rtol=1e-5)
+        self.assertAllCloseAccordingToType(
+            np.array([2.715679168701172, 3.715679168701172]),
+            self.evaluate(var1),
+            float_rtol=1e-5)
+
+  def testTensorLearningRate(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        ada_opt = adagrad.AdagradOptimizer(
+            constant_op.constant(3.0), initial_accumulator_value=0.1)
+        ada_update = ada_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Run 3 steps of adagrad
+        for _ in range(3):
+          ada_update.run()
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-1.6026098728179932, -0.6026098728179932]),
+            self.evaluate(var0),
+            float_rtol=1e-5)
+        self.assertAllCloseAccordingToType(
+            np.array([2.715679168701172, 3.715679168701172]),
+            self.evaluate(var1),
+            float_rtol=1e-5)
+
+  def testSharing(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        ada_opt = adagrad.AdagradOptimizer(3.0)
+        # Apply the optimizer twice.  Both applications will use
+        # the same accums.
+        ada_update1 = ada_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        ada_update2 = ada_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.assertEqual(["accumulator"], ada_opt.get_slot_names())
+        slot0 = ada_opt.get_slot(var0, "accumulator")
+        self.assertEqual(slot0.get_shape(), var0.get_shape())
+        slot1 = ada_opt.get_slot(var1, "accumulator")
+        self.assertEqual(slot1.get_shape(), var1.get_shape())
+        self.evaluate(variables.global_variables_initializer())
+
+        # Fetch params to validate initial values.
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Mix the first and the second adagrad for 3 steps.
+        ada_update1.run()
+        ada_update2.run()
+        ada_update1.run()
+        # Validate updated params (the same as with only 1 Adagrad).
+        self.assertAllCloseAccordingToType(
+            np.array([-1.6026098728179932, -0.6026098728179932]),
+            self.evaluate(var0),
+            float_rtol=1e-5)
+        self.assertAllCloseAccordingToType(
+            np.array([2.715679168701172, 3.715679168701172]),
+            self.evaluate(var1),
+            float_rtol=1e-5)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/adam_test.py

@@ -0,0 +1,185 @@
+# 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.
+# ==============================================================================
+"""Tests for Adam."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adam
+
+
+def adam_update_numpy(param,
+                      g_t,
+                      t,
+                      m,
+                      v,
+                      alpha=0.001,
+                      beta1=0.9,
+                      beta2=0.999,
+                      epsilon=1e-8):
+  alpha_t = alpha * np.sqrt(1 - beta2**t) / (1 - beta1**t)
+
+  m_t = beta1 * m + (1 - beta1) * g_t
+  v_t = beta2 * v + (1 - beta2) * g_t * g_t
+
+  param_t = param - alpha_t * m_t / (np.sqrt(v_t) + epsilon)
+  return param_t, m_t, v_t
+
+
+class AdamOptimizerTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    for dtype in self.float_types | self.complex_types:
+      # TODO: test fails for float16 due to excessive precision requirements.
+      if dtype in [np.float16, dtypes.bfloat16.as_numpy_dtype]:
+        continue
+      with self.session(), self.test_scope():
+        variable_scope.get_variable_scope().set_use_resource(True)
+
+        # Initialize variables for numpy implementation.
+        m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+        var0_np = np.array([1.0, 2.0], dtype=dtype)
+        grads0_np = np.array([0.1, 0.1], dtype=dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype)
+        grads1_np = np.array([0.01, 0.01], dtype=dtype)
+
+        var0 = resource_variable_ops.ResourceVariable(var0_np)
+        var1 = resource_variable_ops.ResourceVariable(var1_np)
+        grads0 = array_ops.placeholder(dtype)
+        grads1 = array_ops.placeholder(dtype)
+        opt = adam.AdamOptimizer()
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        beta1_power, beta2_power = opt._get_beta_accumulators()
+
+        # Run 3 steps of Adam
+        for t in range(1, 4):
+          self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+          self.assertAllCloseAccordingToType(0.999**t,
+                                             self.evaluate(beta2_power))
+          update.run(feed_dict={grads0: grads0_np, grads1: grads1_np})
+
+          var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+          var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+          # Validate updated params
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  def testTensorLearningRate(self):
+    for dtype in self.float_types | self.complex_types:
+      # TODO: test fails for float16 due to excessive precision requirements.
+      if dtype in [np.float16, dtypes.bfloat16.as_numpy_dtype]:
+        continue
+      with self.session(), self.test_scope():
+        variable_scope.get_variable_scope().set_use_resource(True)
+
+        # Initialize variables for numpy implementation.
+        m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+        var0_np = np.array([1.0, 2.0], dtype=dtype)
+        grads0_np = np.array([0.1, 0.1], dtype=dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype)
+        grads1_np = np.array([0.01, 0.01], dtype=dtype)
+
+        var0 = resource_variable_ops.ResourceVariable(var0_np)
+        var1 = resource_variable_ops.ResourceVariable(var1_np)
+        grads0 = array_ops.placeholder(dtype)
+        grads1 = array_ops.placeholder(dtype)
+        opt = adam.AdamOptimizer(constant_op.constant(0.001))
+        update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        beta1_power, beta2_power = opt._get_beta_accumulators()
+
+        # Run 3 steps of Adam
+        for t in range(1, 4):
+          self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+          self.assertAllCloseAccordingToType(0.999**t,
+                                             self.evaluate(beta2_power))
+          update.run(feed_dict={grads0: grads0_np, grads1: grads1_np})
+
+          var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+          var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+          # Validate updated params
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  def testSharing(self):
+    for dtype in self.float_types | self.complex_types:
+      # TODO: test fails for float16 due to excessive precision requirements.
+      if dtype in [np.float16, dtypes.bfloat16.as_numpy_dtype]:
+        continue
+      with self.session(), self.test_scope():
+        variable_scope.get_variable_scope().set_use_resource(True)
+
+        # Initialize variables for numpy implementation.
+        m0, v0, m1, v1 = 0.0, 0.0, 0.0, 0.0
+        var0_np = np.array([1.0, 2.0], dtype=dtype)
+        grads0_np = np.array([0.1, 0.1], dtype=dtype)
+        var1_np = np.array([3.0, 4.0], dtype=dtype)
+        grads1_np = np.array([0.01, 0.01], dtype=dtype)
+
+        var0 = resource_variable_ops.ResourceVariable(var0_np)
+        var1 = resource_variable_ops.ResourceVariable(var1_np)
+        grads0 = array_ops.placeholder(dtype)
+        grads1 = array_ops.placeholder(dtype)
+        opt = adam.AdamOptimizer()
+        update1 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        update2 = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+
+        beta1_power, beta2_power = opt._get_beta_accumulators()
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+        # Run 3 steps of intertwined Adam1 and Adam2.
+        for t in range(1, 4):
+          self.assertAllCloseAccordingToType(0.9**t, self.evaluate(beta1_power))
+          self.assertAllCloseAccordingToType(0.999**t,
+                                             self.evaluate(beta2_power))
+          if t % 2 == 0:
+            update1.run(feed_dict={grads0: grads0_np, grads1: grads1_np})
+          else:
+            update2.run(feed_dict={grads0: grads0_np, grads1: grads1_np})
+
+          var0_np, m0, v0 = adam_update_numpy(var0_np, grads0_np, t, m0, v0)
+          var1_np, m1, v1 = adam_update_numpy(var1_np, grads1_np, t, m1, v1)
+
+          # Validate updated params
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/add_n_test.py

@@ -0,0 +1,80 @@
+# 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 AddN."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import list_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class XlaAddNTest(xla_test.XLATestCase):
+
+  def testAddTensorLists(self):
+    with self.session(), self.test_scope():
+      l1 = list_ops.tensor_list_reserve(
+          element_shape=[], element_dtype=dtypes.float32, num_elements=3)
+      l2 = list_ops.tensor_list_reserve(
+          element_shape=[], element_dtype=dtypes.float32, num_elements=3)
+      l1 = list_ops.tensor_list_set_item(l1, 0, 5.)
+      l2 = list_ops.tensor_list_set_item(l2, 2, 10.)
+
+      l = math_ops.add_n([l1, l2])
+      self.assertAllEqual(
+          list_ops.tensor_list_stack(l, element_dtype=dtypes.float32),
+          [5.0, 0.0, 10.0])
+
+  def testAddTensorListsFailsIfLeadingDimsMismatch(self):
+    with self.session(), self.test_scope():
+      l1 = list_ops.tensor_list_reserve(
+          element_shape=[], element_dtype=dtypes.float32, num_elements=2)
+      l2 = list_ops.tensor_list_reserve(
+          element_shape=[], element_dtype=dtypes.float32, num_elements=3)
+      l = math_ops.add_n([l1, l2])
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          "TensorList arguments to AddN must all have the same shape"):
+        list_ops.tensor_list_stack(l, element_dtype=dtypes.float32).eval()
+
+  def testAddTensorListsFailsIfElementShapesMismatch(self):
+    with self.session() as session, self.test_scope():
+      # Use placeholders instead of constant values for shapes to prevent TF's
+      # shape inference from catching this early.
+      l1_element_shape = array_ops.placeholder(dtype=dtypes.int32)
+      l2_element_shape = array_ops.placeholder(dtype=dtypes.int32)
+      l1 = list_ops.tensor_list_reserve(
+          element_shape=l1_element_shape,
+          element_dtype=dtypes.float32,
+          num_elements=3)
+      l2 = list_ops.tensor_list_reserve(
+          element_shape=l2_element_shape,
+          element_dtype=dtypes.float32,
+          num_elements=3)
+      l = math_ops.add_n([l1, l2])
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          "TensorList arguments to AddN must all have the same shape"):
+        session.run(
+            list_ops.tensor_list_stack(l, element_dtype=dtypes.float32), {
+                l1_element_shape: [],
+                l2_element_shape: [2]
+            })
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/approx_topk_test.py

@@ -0,0 +1,254 @@
+# Copyright 2022 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 approx_max_k and approx_min_k."""
+
+import itertools
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import test
+from tensorflow.python.eager.def_function import function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import variables
+
+
+class ApproxTopkTest(test.TestCase, parameterized.TestCase):
+
+  def setUp(self):
+    test.TestCase.setUp(self)
+    self._rng = np.random.default_rng(42)
+
+  def compute_recall(self, result_neighbors, ground_truth_neighbors):
+    """Computes the recall of an approximate nearest neighbor search.
+
+    Args:
+      result_neighbors: int32 numpy array of the shape [num_queries,
+        neighbors_per_query] where the values are the indices of the dataset.
+      ground_truth_neighbors: int32 numpy array of with shape [num_queries,
+        ground_truth_neighbors_per_query] where the values are the indices of
+        the dataset.
+
+    Returns:
+      The recall.
+    """
+    self.assertLen(result_neighbors.shape, 2)
+    self.assertLen(ground_truth_neighbors.shape, 2)
+    self.assertEqual(result_neighbors.shape[0], ground_truth_neighbors.shape[0])
+    gt_sets = [set(np.asarray(x)) for x in ground_truth_neighbors]
+
+    def hits_per_q(q, nn_per_q):
+      return len(list(x for x in nn_per_q if x.item() in gt_sets[q]))
+
+    hits = sum(
+        hits_per_q(q, nn_per_q) for q, nn_per_q in enumerate(result_neighbors))
+    return hits / ground_truth_neighbors.size
+
+  @parameterized.parameters(
+      itertools.product(
+          [1, 10],  # k
+          [100, 500],  # row_size
+          [1, 10, 128],  # num_rows
+          [True, False],  # aggregate_to_topk
+      ))
+  def test_non_fused_max_k(self, k, row_size, num_rows, aggregate_to_topk):
+    row = np.arange(row_size, dtype=np.float32)
+    db = np.stack(list(self._rng.permutation(row) for _ in range(num_rows)))
+
+    @function(jit_compile=True)
+    def ann(db, k):
+      return nn_ops.approx_max_k(db, k, aggregate_to_topk=aggregate_to_topk)
+
+    with ops.device('/device:TPU:0'):
+      db_op = variables.Variable(db)
+      result = ann(db_op, k)[1]
+
+    gt = np.argsort(-db)[:, :k]
+    ann_recall = self.compute_recall(result.numpy(), gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+  @parameterized.parameters(
+      itertools.product(
+          [1, 10],  # k
+          [100, 500],  # row_size
+          [1, 10, 128],  # num_rows
+          [True, False],  # aggregate_to_topk
+      ))
+  def test_non_fused_min_k(self, k, row_size, num_rows, aggregate_to_topk):
+    # Use the new rng api
+    row = np.arange(row_size, dtype=np.float32)
+    db = np.stack(list(self._rng.permutation(row) for _ in range(num_rows)))
+
+    @function(jit_compile=True)
+    def ann(db, k=10):
+      return nn_ops.approx_min_k(db, k, aggregate_to_topk=aggregate_to_topk)
+
+    with ops.device('/device:TPU:0'):
+      db_op = variables.Variable(db)
+      result = ann(db_op, k)[1]
+
+    gt = np.argsort(db)[:, :k]
+    ann_recall = self.compute_recall(result.numpy(), gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+  @parameterized.parameters(
+      itertools.product(
+          [1, 10],  # k
+          [100, 500],  # db_size
+          [1, 10, 128],  # qy_size
+          [2, 32],  # feature dim
+      ))
+  # MIPS = Maximal Inner Product Search
+  def test_mips(self, k, db_size, qy_size, feature_dim):
+    qy = self._rng.random([qy_size, feature_dim], dtype=np.float32)
+    db = self._rng.random([db_size, feature_dim], dtype=np.float32)
+
+    @function(jit_compile=True)
+    def ann(qy, db, k):
+      scores = math_ops.matmul(qy, db, transpose_b=True)
+      return nn_ops.approx_max_k(scores, k)
+
+    with ops.device('/device:TPU:0'):
+      qy_op = variables.Variable(qy)
+      db_op = variables.Variable(db)
+      result = ann(qy_op, db_op, k)[1]
+      scores = -math_ops.matmul(qy_op, db_op, transpose_b=True)
+
+    gt = np.argsort(scores.numpy())[:, :k]
+    ann_recall = self.compute_recall(result.numpy(), gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+  @parameterized.parameters(
+      itertools.product(
+          [1, 10],  # k
+          [100, 500],  # db_size
+          [10, 128],  # qy_size
+          [2, 8],  # feature dim
+      ))
+  # L2ANN = Approximate Nearest Neighbor search in the L2 metric space
+  def test_l2ann(self, k, db_size, qy_size, feature_dim):
+    qy = self._rng.random([qy_size, feature_dim], dtype=np.float32)
+    db = self._rng.random([db_size, feature_dim], dtype=np.float32)
+    db_half_norm_sq = np.linalg.norm(db, axis=1)**2 / 2
+
+    @function(jit_compile=True)
+    def ann(qy, db, db_half_norm_sq, k):
+      scores = db_half_norm_sq - math_ops.matmul(qy, db, transpose_b=True)
+      return nn_ops.approx_min_k(scores, k)
+
+    with ops.device('/device:TPU:0'):
+      qy_op = variables.Variable(qy)
+      db_op = variables.Variable(db)
+      db_half_norm_sq_op = variables.Variable(db_half_norm_sq)
+      result = ann(qy_op, db_op, db_half_norm_sq_op, k)[1]
+      scores = db_half_norm_sq_op - math_ops.matmul(
+          qy_op, db_op, transpose_b=True)
+
+    gt = np.argsort(scores.numpy())[:, :k]
+    ann_recall = self.compute_recall(result.numpy(), gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+  def test_highdim(self):
+    db = self._rng.random([2, 10, 200, 3], dtype=np.float32)
+    k = 5
+
+    @function(jit_compile=True)
+    def ann(db, k):
+      return nn_ops.approx_min_k(db, k=k, reduction_dimension=2)
+
+    with ops.device('/device:TPU:0'):
+      db_op = variables.Variable(db)
+      result = ann(db_op, k)[1]
+
+    gt = np.argsort(db, axis=2)[:, :, :k, :]
+    flat_idx = np.reshape(
+        np.transpose(result.numpy(), [0, 1, 3, 2]), [2 * 10 * 3, k])
+    flat_gt = np.reshape(np.transpose(gt, [0, 1, 3, 2]), [2 * 10 * 3, k])
+    ann_recall = self.compute_recall(flat_idx, flat_gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+  @parameterized.parameters(
+      itertools.product(
+          [dtypes.bfloat16, dtypes.float16, dtypes.float32],
+          [1, 10],  # k
+          [100, 500],  # row_size
+          [1, 10, 128],  # num_rows
+      ))
+  def test_gradients(self, dtype, k, row_size, num_rows):
+    row = np.arange(row_size, dtype=np.float32)
+    db = np.stack(list(self._rng.permutation(row) for _ in range(num_rows)))
+    out_grads = self._rng.random([num_rows, k])
+
+    @function(jit_compile=True)
+    def ann_with_grads(db, out_grads):
+      with backprop.GradientTape() as tape:
+        val, idx = nn_ops.approx_max_k(db, k)
+      result_in_grads = tape.gradient(val, db, out_grads)
+      lifted_k_idx = array_ops.reshape(idx, [num_rows, k, 1])
+      iota_idx = array_ops.broadcast_to(
+          array_ops.reshape(math_ops.range(num_rows), [num_rows, 1, 1]),
+          [num_rows, k, 1])
+      lifted_idx = array_ops.concat([iota_idx, lifted_k_idx], axis=2)
+      k_idx_s = array_ops.reshape(lifted_idx, [num_rows * k, 2])
+      k_gra_s = array_ops.reshape(out_grads, [num_rows * k])
+      expected_in_grads = array_ops.scatter_nd(k_idx_s, k_gra_s,
+                                               [num_rows, row_size])
+      return [expected_in_grads, result_in_grads]
+
+    with ops.device('/device:TPU:0'):
+      db_op = variables.Variable(db, dtype=dtype)
+      out_grads_op = variables.Variable(out_grads, dtype=dtype)
+      expected_in_grads, result_in_grads = ann_with_grads(db_op, out_grads_op)
+
+    self.assertAllClose(expected_in_grads, result_in_grads)
+
+  # Tests that multiple ops are supported and the comparison functions are
+  # renamed properly to avoid conflict while using the MLIR bridge.
+  def test_multiple_ops(self):
+    k = 1
+
+    row_size = 100
+    num_rows = 10
+
+    row = np.arange(row_size, dtype=np.float32)
+    db1 = np.stack(list(self._rng.permutation(row) for _ in range(num_rows)))
+    db2 = np.stack(list(self._rng.permutation(row) for _ in range(num_rows)))
+
+    @function(jit_compile=True)
+    def ann(db1, db2):
+      result1 = nn_ops.approx_max_k(db1, k, aggregate_to_topk=True)
+      result2 = nn_ops.approx_max_k(db2, k, aggregate_to_topk=True)
+      return (result1, result2)
+
+    with ops.device('/device:TPU:0'):
+      db1_op = variables.Variable(db1)
+      db2_op = variables.Variable(db2)
+      result1, result2 = ann(db1_op, db2_op)
+
+    gt = np.argsort(-db1)[:, :k]
+    ann_recall = self.compute_recall(result1[1].numpy(), gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+    gt = np.argsort(-db2)[:, :k]
+    ann_recall = self.compute_recall(result2[1].numpy(), gt)
+    self.assertGreaterEqual(ann_recall, 0.95)
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/tests/argminmax_test.py

@@ -0,0 +1,95 @@
+# 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 ArgMin and ArgMax Ops."""
+
+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 test
+
+
+class ArgMinMaxTest(xla_test.XLATestCase):
+
+  def _assertOpOutputMatchesExpected(self, op, axis, output_type, op_input,
+                                     expected):
+    """Verifies that 'op' produces 'expected' when fed input 'op_input' .
+
+    Args:
+      op: argmin or argmax operator to test.
+      axis: integer axis to reduce across.
+      output_type: numpy datatype of the output to produce.
+      op_input: numpy input array to use as input to 'op'.
+      expected: numpy array representing the expected output of 'op'.
+    """
+    with self.session() as session:
+      with self.test_scope():
+        pinp = array_ops.placeholder(
+            dtypes.as_dtype(op_input.dtype), op_input.shape, name="a")
+        output = op(pinp, axis=axis, output_type=output_type)
+      result = session.run(output, {pinp: op_input})
+      self.assertAllEqual(result, expected)
+
+  def testArgMinMax(self):
+    # Complex numbers do not support argmin/argmax.
+    minmax_types = self.all_types & {np.int32, np.int64}
+    for dtype in self.int_types | self.float_types:
+      # output_type is a numpy data type that is used to specify the desired
+      # output type of the op as well as to convert the Python number to the
+      # array scalar of the type.
+      for output_type in minmax_types:
+        self._assertOpOutputMatchesExpected(
+            math_ops.argmax,
+            axis=0,
+            output_type=output_type,
+            op_input=np.array([1, 10, 27, 3, 3, 4], dtype=dtype),
+            expected=output_type(2))
+        self._assertOpOutputMatchesExpected(
+            math_ops.argmax,
+            axis=0,
+            output_type=output_type,
+            op_input=np.array([[4, 1, 7], [3, 2, 4]], dtype=dtype),
+            expected=np.array([0, 1, 0], dtype=output_type))
+        self._assertOpOutputMatchesExpected(
+            math_ops.argmax,
+            axis=1,
+            output_type=output_type,
+            op_input=np.array([[4, 1], [3, 2]], dtype=dtype),
+            expected=np.array([0, 0], dtype=output_type))
+
+        self._assertOpOutputMatchesExpected(
+            math_ops.argmin,
+            axis=0,
+            output_type=output_type,
+            op_input=np.array([3, 10, 27, 3, 2, 4], dtype=dtype),
+            expected=output_type(4))
+        self._assertOpOutputMatchesExpected(
+            math_ops.argmin,
+            axis=0,
+            output_type=output_type,
+            op_input=np.array([[4, 1, 7], [3, 2, 4]], dtype=dtype),
+            expected=np.array([1, 0, 1], dtype=output_type))
+        self._assertOpOutputMatchesExpected(
+            math_ops.argmin,
+            axis=1,
+            output_type=output_type,
+            op_input=np.array([[4, 1], [3, 2]], dtype=dtype),
+            expected=np.array([1, 1], dtype=output_type))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/async_comp_test.py

@@ -0,0 +1,95 @@
+# 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.
+# ==============================================================================
+"""Tests for asynchronous compilation on the CPU and GPU devices."""
+
+import os
+
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def RunMetadataLabels(run_metadata):
+  """Returns all labels in run_metadata."""
+  labels = []
+  for dev_stats in run_metadata.step_stats.dev_stats:
+    for node_stats in dev_stats.node_stats:
+      labels.append(node_stats.timeline_label)
+  return labels
+
+
+def InLabels(labels, substr):
+  """Returns true iff one of the labels contains substr."""
+  return any(substr in x for x in labels)
+
+
+def MetadataHasXlaRunOp(run_metadata):
+  """Returns true if there are XlaRun kernels in run_metadata's timeline."""
+
+  # TODO(phawkins): find a less hacky way to test whether a kernel ran.
+  return InLabels(RunMetadataLabels(run_metadata), "_XlaRun")
+
+
+class AsyncCompilationTest(test.TestCase):
+
+  # Asynchrobnous compilation uses the existing fallback path and existing
+  # compiler. This test only tests that asynchronous compilation is performed.
+  def testAsyncCompilationJit(self):
+
+    @function.Defun(compiled=True)
+    def CompiledFunction(x):
+      return math_ops.log(x)
+
+    with session_lib.Session() as sess:
+      x = array_ops.placeholder(dtypes.float32)
+      y = CompiledFunction(x)
+
+      run_metadata = config_pb2.RunMetadata()
+      sess.run(
+          y,
+          feed_dict={x: [0.] * 60},
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      # For The first iteration, the fall back path is chosen.
+      hasXlaRunOp = MetadataHasXlaRunOp(run_metadata)
+      self.assertFalse(hasXlaRunOp)
+
+      # Execute the session until after asynchronous compilation is finished
+      # and the compiled cluster has been executed once.
+      while (not hasXlaRunOp):
+        run_metadata = config_pb2.RunMetadata()
+        sess.run(
+            y,
+            feed_dict={x: [0.] * 60},
+            run_metadata=run_metadata,
+            options=config_pb2.RunOptions(
+                trace_level=config_pb2.RunOptions.FULL_TRACE))
+        hasXlaRunOp = MetadataHasXlaRunOp(run_metadata)
+
+
+if __name__ == "__main__":
+  os.environ["TF_XLA_FLAGS"] = ("--tf_xla_async_compilation=true " +
+                                "--tf_xla_enable_lazy_compilation=true " +
+                                os.environ.get("TF_XLA_FLAGS", ""))
+  # This test is using Tensorflow sessions which are not compatible with eager
+  # mode.
+  ops.disable_eager_execution()
+  test.main()

tensorflow/compiler/tests/bincount_op_test.py

@@ -0,0 +1,74 @@
+# Copyright 2023 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 bincount using the XLA JIT."""
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.platform import googletest
+
+
+class BincountTest(xla_test.XLATestCase):
+
+  def testInputRank0(self):
+    with self.session():
+      with self.test_scope():
+        bincount = gen_math_ops.bincount(arr=6, size=804, weights=[52, 351])
+
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          (
+              "`weights` must be the same shape as `arr` or a length-0"
+              " `Tensor`, in which case it acts as all weights equal to 1."
+          ),
+      ):
+        self.evaluate(bincount)
+
+  def testNegativeInputConstant(self):
+    with self.session():
+      with self.test_scope():
+
+        @def_function.function(jit_compile=True)
+        def f():
+          return gen_math_ops.dense_bincount(
+              input=array_ops.constant([0, -1, 2]), size=3, weights=[]
+          )
+
+        with self.assertRaisesRegex(
+            errors.InvalidArgumentError, "Input arr must be non-negative!"
+        ):
+          self.evaluate(f())
+
+  def testNegativeInputConstantBincount(self):
+    with self.session():
+      with self.test_scope():
+
+        @def_function.function(jit_compile=True)
+        def f():
+          return gen_math_ops.bincount(
+              arr=array_ops.constant([0, -1, 2]), size=3, weights=[]
+          )
+
+        with self.assertRaisesRegex(
+            errors.InvalidArgumentError, "Input arr must be non-negative!"
+        ):
+          self.evaluate(f())
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  googletest.main()

tensorflow/compiler/tests/bucketize_op_test.py

@@ -0,0 +1,76 @@
+# Copyright 2016 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 bucketize_op."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class BucketizationOpTest(xla_test.XLATestCase):
+
+  def testInt(self):
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.int32)
+      with self.test_scope():
+        op = math_ops._bucketize(p, boundaries=[0, 3, 8, 11])
+      expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4]
+      self.assertAllEqual(expected_out,
+                          sess.run(op, {p: [-5, 0, 2, 3, 5, 8, 10, 11, 12]}))
+
+  def testFloat(self):
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.float32)
+      with self.test_scope():
+        op = math_ops._bucketize(p, boundaries=[0., 3., 8., 11.])
+      expected_out = [0, 1, 1, 2, 2, 3, 3, 4, 4]
+      self.assertAllEqual(
+          expected_out,
+          sess.run(op, {p: [-5., 0., 2., 3., 5., 8., 10., 11., 12.]}))
+
+  def test2DInput(self):
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.float32)
+      with self.test_scope():
+        op = math_ops._bucketize(p, boundaries=[0, 3, 8, 11])
+      expected_out = [[0, 1, 1, 2, 2], [3, 3, 4, 4, 1]]
+      self.assertAllEqual(
+          expected_out, sess.run(op,
+                                 {p: [[-5, 0, 2, 3, 5], [8, 10, 11, 12, 0]]}))
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testInvalidBoundariesOrder(self):
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.int32)
+      with self.test_scope():
+        op = math_ops._bucketize(p, boundaries=[0, 8, 3, 11])
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  "Expected sorted boundaries"):
+        sess.run(op, {p: [-5, 0]})
+
+  def testBoundariesNotList(self):
+    with self.session():
+      with self.assertRaisesRegex(TypeError, "Expected list.*"):
+        p = array_ops.placeholder(dtypes.int32)
+        with self.test_scope():
+          math_ops._bucketize(p, boundaries=0)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/build_combined_defs.bzl

@@ -0,0 +1,63 @@
+"""Build rule for combining Tensorflow/XLA tests."""
+
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow/compiler/tests:build_defs.bzl", "tf_xla_py_test")
+
+def parse_label_name(label):
+    """Parse a label into just the name.
+
+    Args:
+      label: string in relative or absolute form.
+
+    Returns:
+      The name of the label.
+    """
+    colon_split = label.split(":")
+    if len(colon_split) == 1:  # no ":" in label
+        return label
+    return colon_split[-1]
+
+def tf_xla_combined_py_test(name = "", package = None, tests = [], **kwargs):
+    """Generates combined tf_xla_py_test targets, one per XLA backend.
+
+    All srcs found in the list tests are combined into one new test which is then passed on to
+    tf_xla_py_test which creates a new target per XLA backend.
+
+    Args:
+      name: Name of the target.
+      package: The package that all tests in tests belong to.
+      tests: The test targets to be combined and tested. Assumes all tests are in the same package.
+      **kwargs: keyword arguments passed onto the tf_xla_py_test rule.
+    """
+
+    test_file = name + ".py"
+
+    # run the generator to create the combined test file containing all the tests in test_files
+    # redirecting the output of the generator to test_file.
+    native.genrule(
+        name = name + "_gen",
+        testonly = 1,
+        srcs = tests,
+        outs = [test_file],
+        cmd = """
+mkdir -p $(@D) && cat > $@ << EOF
+from tensorflow.python.platform import test
+%s
+
+if __name__ == "__main__":
+  test.main()
+EOF
+        """ % "\n".join(["from %s.%s import *" % (package, parse_label_name(test)[:-4]) for test in tests]),
+        tools = [],
+        tags = ["generated_python_test=%s.%s" % (package, name)],
+    )
+
+    tf_xla_py_test(
+        name = name,
+        test_rule = py_test,
+        srcs = [test_file],
+        deps = [
+            "//tensorflow/python/platform:client_testlib",
+        ] + tests,
+        **kwargs
+    )

tensorflow/compiler/tests/build_defs.bzl

@@ -0,0 +1,168 @@
+"""Build rules for Tensorflow/XLA testing."""
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow/compiler/tests:plugin.bzl", "plugins")
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "tf_cuda_tests_tags",
+    "tf_exec_properties",
+)
+
+all_backends = ["cpu", "gpu"] + plugins.keys()
+
+def tf_xla_py_test(
+        name,
+        srcs = [],
+        deps = [],
+        tags = [],
+        data = [],
+        main = None,
+        enabled_backends = None,
+        disabled_backends = None,
+        use_xla_device = True,
+        enable_mlir_bridge = True,
+        test_rule = py_test,
+        **kwargs):
+    """Generates py_test targets, one per XLA backend.
+
+    This rule generates py_test() targets named name_backend, for each backend
+    in all_backends. The rule also generates a test suite with named `name` that
+    tests all backends for the test.
+
+    For example, the following rule generates test cases foo_test_cpu,
+    foo_test_gpu, and a test suite name foo_test that tests both.
+    tf_xla_py_test(
+        name="foo_test",
+        srcs="foo_test.py",
+        deps=[...],
+    )
+
+    Args:
+      name: Name of the target.
+      srcs: Sources for the target.
+      deps: Dependencies of the target.
+      tags: Tags to apply to the generated targets.
+      data: Data dependencies of the target.
+      main: Same as py_test's main attribute.
+      enabled_backends: A list of backends that should be tested. Supported
+        values include "cpu" and "gpu". If not specified, defaults to None.
+      disabled_backends: A list of backends that should not be tested. Supported
+        values include "cpu" and "gpu". If not specified, defaults to None.
+      use_xla_device: If true then the --test_device argument is set to XLA_CPU
+        and XLA_GPU for the CPU and GPU tests.  Otherwise it is set to CPU and
+        GPU.
+      enable_mlir_bridge: If true, then runs the test with and without mlir
+        bridge enabled.
+      **kwargs: keyword arguments passed onto the generated py_test() rules.
+    """
+    if enabled_backends == None:
+        enabled_backends = all_backends
+    if disabled_backends == None:
+        disabled_backends = []
+    if type(disabled_backends) != "list":
+        fail("disabled_backends must be a list of strings", "disabled_backends")
+
+    backends = [b for b in enabled_backends if b not in disabled_backends]
+    test_names = []
+
+    if use_xla_device:
+        cpu_xla_device = "XLA_CPU"
+        gpu_xla_device = "XLA_GPU"
+    else:
+        cpu_xla_device = "CPU"
+        gpu_xla_device = "GPU"
+
+    py_library(
+        name = name + "_lib",
+        srcs = srcs,
+        deps = deps,
+        testonly = 1,
+    )
+    for backend in backends:
+        test_name = "{}_{}".format(name, backend)
+        backend_tags = ["tf_xla_{}".format(backend)]
+        backend_args = []
+        backend_deps = []
+        backend_data = []
+        if backend == "cpu":
+            backend_args += [
+                "--test_device=" + cpu_xla_device,
+                "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_UINT8,DT_QUINT8,DT_INT8,DT_QINT8,DT_INT32,DT_QINT32,DT_INT64,DT_BOOL,DT_COMPLEX64,DT_COMPLEX128",
+            ]
+        elif backend in ("gpu", "gpu_a100", "gpu_h100"):
+            backend_args += [
+                "--test_device=" + gpu_xla_device,
+                "--types=DT_HALF,DT_FLOAT,DT_DOUBLE,DT_UINT8,DT_QUINT8,DT_INT8,DT_QINT8,DT_INT32,DT_QINT32,DT_INT64,DT_BOOL,DT_COMPLEX64,DT_COMPLEX128,DT_BFLOAT16",
+            ]
+            backend_tags += tf_cuda_tests_tags()
+        elif backend in plugins:
+            backend_args += [
+                "--test_device=" + plugins[backend]["device"],
+                "--types=" + plugins[backend]["types"],
+            ]
+            backend_tags += plugins[backend]["tags"]
+            backend_args += plugins[backend]["args"]
+            backend_deps += plugins[backend]["deps"]
+            backend_data += plugins[backend]["data"]
+        else:
+            fail("Unknown backend {}".format(backend))
+
+        test_tags = tags + backend_tags
+
+        enable_mlir_bridge_options = [False]
+        if enable_mlir_bridge:
+            enable_mlir_bridge_options.append(True)
+
+        for mlir_option in enable_mlir_bridge_options:
+            extra_dep = []
+            extra_tag = []
+            updated_name = test_name
+
+            mlir_bridge_dep = "//tensorflow/python/framework:is_mlir_bridge_test_true"
+            has_mlir_dep = (mlir_bridge_dep in deps)
+            if mlir_option:
+                if updated_name.endswith("_test"):
+                    updated_name = updated_name[:-5]
+                updated_name += "_mlir_bridge_test"
+                extra_dep = [] if has_mlir_dep else [mlir_bridge_dep]
+
+                # Mark gpu mlir_bridge tests as ondemand
+                #
+                # This is for testing book keeping because the bridge does not have any gpu specific
+                # logic at this time, so CPU testing is good enough and cheaper.
+                extra_tag = ["ondemand"] if backend in ("gpu", "gpu_a100", "gpu_h100") else []
+            elif has_mlir_dep:
+                # Some tests run only with mlir_bridge by explicitly adding the MLIR
+                # bridge dep so if the dep is already present skip non MLIR
+                # version.
+                continue
+
+            # Rules may set exec_properties, but Google has internal
+            # exec_properties values so they don't merge easily. Just strip them
+            # all for now.
+            kwargs.pop("exec_properties", {})
+            test_rule(
+                name = updated_name,
+                srcs = srcs,
+                srcs_version = "PY3",
+                args = backend_args,
+                main = "{}.py".format(name) if main == None else main,
+                data = data + backend_data,
+                deps = deps + backend_deps + extra_dep + [name + "_lib"],
+                tags = test_tags + extra_tag,
+                exec_properties = tf_exec_properties({"tags": test_tags}),
+                **kwargs
+            )
+            test_names.append(updated_name)
+    native.test_suite(name = name, tests = test_names)
+
+def tf_xla_py_strict_test(**kwargs):
+    tf_xla_py_test(**kwargs)
+
+def generate_backend_suites(backends = []):
+    """Generates per-backend test_suites that run all tests for a backend."""
+    if not backends:
+        backends = all_backends
+    for backend in backends:
+        native.test_suite(name = "%s_tests" % backend, tags = ["tf_xla_%s" % backend])

tensorflow/compiler/tests/case_test.py

@@ -0,0 +1,83 @@
+# 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 case statements in XLA."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_switch_case
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import test
+
+
+class CaseTest(xla_test.XLATestCase):
+
+  def testCaseBasic(self):
+
+    @def_function.function(jit_compile=True)
+    def switch_case_test(branch_index):
+
+      def f1():
+        return array_ops.constant(17)
+
+      def f2():
+        return array_ops.constant(31)
+
+      def f3():
+        return array_ops.constant(-1)
+
+      return control_flow_switch_case.switch_case(
+          branch_index, branch_fns={
+              0: f1,
+              1: f2
+          }, default=f3)
+
+    with ops.device(self.device):
+      self.assertEqual(switch_case_test(array_ops.constant(0)).numpy(), 17)
+      self.assertEqual(switch_case_test(array_ops.constant(1)).numpy(), 31)
+      self.assertEqual(switch_case_test(array_ops.constant(2)).numpy(), -1)
+      self.assertEqual(switch_case_test(array_ops.constant(3)).numpy(), -1)
+
+  def testBranchIsPruned(self):
+
+    @def_function.function(jit_compile=True)
+    def switch_case_test():
+      branch_index = array_ops.constant(0)
+
+      def f1():
+        return array_ops.constant(17)
+
+      def f2():
+        # Some operations that XLA cannot compile.
+        image_ops.decode_image(io_ops.read_file('/tmp/bmp'))
+        return array_ops.constant(31)
+
+      # This tests that we do not try to compile all branches if the branch
+      # index in trivially constant.
+      return control_flow_switch_case.switch_case(
+          branch_index, branch_fns={
+              0: f1,
+              1: f2
+          }, default=f2)
+
+    with ops.device(self.device):
+      self.assertEqual(switch_case_test().numpy(), 17)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/compiler/tests/cast_ops_test.py

@@ -0,0 +1,58 @@
+# 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 lowering of tf.bitcast"""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import io_ops
+from tensorflow.python.platform import test
+
+
+class CastOpsTest(xla_test.XLATestCase):
+
+  def testBitcastToLarger(self):
+    with ops.device('device:{}:0'.format(self.device)):
+
+      def f(x):
+        t = array_ops.bitcast(x, dtypes.float32)
+        return math_ops.reduce_sum(t, axis=1)
+
+      compiled_f = def_function.function(f, jit_compile=True)
+
+      x = random_ops.random_normal([10, 10, 2], dtype=dtypes.float16)
+      with ops.device(self.device):
+        out = f(x)
+        compiled_out = compiled_f(x)
+        self.assertAllClose(out, compiled_out)
+        # 10,10,2--(bitcast-convert)-->10,10--(reduce)-->10
+        self.assertEqual(out.shape[0], 10)
+
+      hlo = compiled_f.experimental_get_compiler_ir(x)(stage='hlo')
+      self.assertIn('f32[10,10]{1,0} bitcast-convert(f16[10,10,2]{2,1,0}', hlo)
+
+  def testBitcastToSmaller(self):
+    pass
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/compiler/tests/cast_test.py

@@ -0,0 +1,126 @@
+# Copyright 2025 The OpenXLA Authors.
+#
+# 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.
+# ==============================================================================
+import platform
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class CastTest(xla_test.XLATestCase):
+
+  def test_cast(self):
+    types = {
+        dtypes.bool,
+        dtypes.float32,
+        dtypes.float64,
+        dtypes.complex64,
+        dtypes.int32,
+        dtypes.int64,
+        dtypes.uint32,
+        dtypes.uint64,
+    }
+    with self.session() as session:
+      for src_type in types:
+        for dst_type in types:
+          self._test_cast(src_type, dst_type, session)
+
+  def test_cast_fp8(self):
+    if platform.system() == "Darwin":
+      # TODO(b/271327511): Fix issue where casts to FP8 very rarely result in
+      # NaN on Mac
+      self.skipTest("Casts to FP8 sometimes result in NaN on Mac")
+    fp8_types = {
+        dtypes.float8_e5m2,
+        dtypes.float8_e4m3fn,
+        dtypes.float8_e4m3fnuz,
+        dtypes.float8_e4m3b11fnuz,
+        dtypes.float8_e5m2fnuz,
+    }
+    other_types = {
+        dtypes.bool,
+        dtypes.float32,
+        dtypes.float64,
+        dtypes.complex64,
+        dtypes.int32,
+        dtypes.int64,
+        dtypes.uint32,
+        dtypes.uint64,
+    }
+    with self.session() as session:
+      for fp8_type in fp8_types:
+        for other_type in other_types | fp8_types:
+          self._test_cast(fp8_type, other_type, session)
+          self._test_cast(other_type, fp8_type, session)
+
+  def _test_cast(self, src_type, dst_type, session):
+    with self.subTest(src_type=src_type, dst_type=dst_type):
+      shapes = [[], [4], [2, 3], [2, 0, 4]]
+      src_np_dtype = src_type.as_numpy_dtype
+      dst_np_dtype = dst_type.as_numpy_dtype
+
+      for shape in shapes:
+        src = np.arange(np.prod(shape)).astype(src_np_dtype)
+
+        if src_type in self.complex_tf_types:
+          src += (np.arange(np.prod(shape)) * 2j).astype(src_np_dtype)
+        src = src.reshape(shape)
+        dst = src.astype(dst_np_dtype)
+        self.assert_op_output_matches_expected(
+            lambda x, dst_type=dst_type: math_ops.cast(x, dst_type),
+            src,
+            expected=dst,
+            local_session=session,
+        )
+
+      # Check special values.
+      if src_type.is_integer:
+        imin = np.iinfo(src_np_dtype).min
+        imax = np.iinfo(src_np_dtype).max
+        if src_type.is_unsigned:
+          src = np.array([imin, imax, 0, 1], dtype=src_np_dtype)
+        else:
+          src = np.array([imin, imax, 0, 1, -1], dtype=src_np_dtype)
+      elif src_type in self.float_tf_types:
+        if dst_type.is_integer:
+          imin = np.iinfo(dst_np_dtype).min
+          imax = np.iinfo(dst_np_dtype).max // 2
+          src = np.array([imin, imax, 0, 1], dtype=src_np_dtype)
+        elif dst_type in self.float_tf_types:
+          fmin = np.finfo(dst_np_dtype).min
+          fmax = np.finfo(dst_np_dtype).max
+          tiny = np.finfo(dst_np_dtype).tiny
+          eps = np.finfo(dst_np_dtype).eps
+          src = np.array(
+              [fmin, fmax, np.nan, eps, -eps, tiny, -tiny, np.inf, -np.inf],
+              dtype=src_np_dtype,
+          )
+      dst = src.astype(dst_np_dtype)
+      self.assert_op_output_matches_expected(
+          lambda x, dst_type=dst_type: math_ops.cast(x, dst_type),
+          src,
+          expected=dst,
+          local_session=session,
+      )
+
+  def test_give_me_a_name(self):
+    pass
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/categorical_op_test.py

@@ -0,0 +1,195 @@
+# 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.
+# ==============================================================================
+"""Tests for multinomial generation ops in the XLA JIT compiler."""
+
+import collections
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import random_seed
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import stateless_random_ops
+from tensorflow.python.platform import googletest
+from tensorflow.python.platform import test
+
+
+# TODO(srvasude): Merge this with
+# third_party/tensorflow/python/kernel_tests/random/multinomial_op_test.py.
+class CategoricalTest(xla_test.XLATestCase):
+  """Test cases for random-number generating operators."""
+
+  def output_dtypes(self):
+    return set(self.int_types).intersection([np.int32, np.int64])
+
+  def _chi2(self, expected, actual):
+    """Returns Chi2 GOF statistic."""
+    actual = np.asarray(actual)
+    expected = np.asarray(expected)
+    diff = actual - expected
+    chi2 = np.sum(diff * diff / expected)
+    return chi2
+
+  def _do_sampling(self, logits, num_samples):
+    """Categorical samples from given input.
+
+    Args:
+      logits: Numpy ndarray of shape [batch_size, num_classes].
+      num_samples: Int; number of samples to draw.
+
+    Returns:
+      Frequencies from sampled classes; shape [batch_size, num_classes].
+    """
+    with self.session(), self.test_scope():
+      random_seed.set_random_seed(1618)
+      op = random_ops.multinomial(logits, num_samples,
+                                  output_dtype=dtypes.int32)
+      d = self.evaluate(op)
+
+    batch_size, num_classes = logits.shape
+    freqs_mat = []
+    for i in range(batch_size):
+      cnts = dict(collections.Counter(d[i, :]))
+
+      # Requires drawn class labels be in range.
+      self.assertLess(max(cnts.keys()), num_classes)
+      self.assertGreaterEqual(min(cnts.keys()), 0)
+
+      freqs = [(cnts[k] * 1. / num_samples if k in cnts else 0)
+               for k in range(num_classes)]
+      freqs_mat.append(freqs)
+
+    return freqs_mat
+
+  def _testRngIsNotConstant(self, rng, dtype, output_dtype):
+    # Tests that 'rng' does not always return the same value.
+    with self.session():
+      with self.test_scope():
+        x = rng(dtype, output_dtype)
+
+      # The random-number generator, if working correctly, should produce the
+      # same output multiple times with low probability.
+      y = self.evaluate(x)
+      z = self.evaluate(x)
+      w = self.evaluate(x)
+
+      # We use exact equality here. If the random-number generator is producing
+      # deterministic output, all three outputs will be bitwise identical.
+      self.assertTrue((not np.array_equal(y, z)) or
+                      (not np.array_equal(z, w)) or
+                      (not np.array_equal(y, w)))
+
+  def testCategoricalIsNotConstant(self):
+    def rng(dtype, output_dtype):
+      return random_ops.multinomial(np.array([[1., 1., 1.]], dtype=dtype), 10,
+                                    output_dtype=output_dtype)
+
+    dtype = np.float32
+    for output_dtype in self.output_dtypes():
+      self._testRngIsNotConstant(rng, dtype, output_dtype)
+
+  @test.disable_with_predicate(
+      pred=test.is_built_with_rocm, skip_message="Test fails on ROCm.")
+  def testCategoricalIsInRange(self):
+    for dtype in self.float_types:
+      for output_dtype in self.output_dtypes():
+        with self.session():
+          with self.test_scope():
+            x = random_ops.multinomial(
+                array_ops.ones(shape=[1, 20], dtype=dtype), 1000,
+                output_dtype=output_dtype)
+          y = self.evaluate(x)
+          self.assertTrue((y >= 0).sum() == 1000)
+          self.assertTrue((y < 20).sum() == 1000)
+
+  def testSamplingCorrectness(self):
+    np.random.seed(1618)  # Make it reproducible.
+    num_samples = 40000
+
+    rand_probs = np.random.dirichlet([1., 1., 2., 3.])
+    rand_probs2 = np.random.dirichlet([1., 4., 5.], size=3)  # batched
+    for probs in [[.5, .5], [.85, .05, .1], rand_probs, rand_probs2]:
+      probs = np.asarray(probs)
+      if len(probs.shape) == 1:
+        probs = probs.reshape(1, probs.size)  # singleton batch
+
+      logits = np.log(probs).astype(np.float32)
+      freqs = self._do_sampling(logits, num_samples)
+
+      # the test here is similar to
+      # python/kernel_tests/random/multinomial_op_test.py
+      # Note that df >= 1 in all these cases. Choosing a cutoff of 1e-3
+      # corresponds to an alpha value of 2.5% for df = 1, and smaller for larger
+      # df.
+      chi2 = self._chi2(probs, freqs)
+      self.assertLess(chi2, 1e-3)
+
+  def testStatelessMultinomialIsInRange(self):
+    for dtype in self.float_types.intersection(
+        [dtypes.float32, dtypes.bfloat16]):
+      for output_dtype in self.output_dtypes():
+        with self.session() as sess:
+          with self.test_scope():
+            seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
+            x = stateless_random_ops.stateless_multinomial(
+                array_ops.ones(shape=[1, 20], dtype=dtype),
+                1000,
+                seed_t,
+                output_dtype=output_dtype)
+          y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]})
+          self.assertTrue((y >= 0).sum() == 1000)
+          self.assertTrue((y < 20).sum() == 1000)
+
+  def testDeterminismMultinomial(self):
+    # Stateless values should be equal iff the seeds are equal (roughly)
+    num_samples = 10
+    with self.session(), self.test_scope():
+      seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
+      seeds = [(x, y) for x in range(5) for y in range(5)] * 3
+      for logits in ([[0.1, 0.25, 0.5, 0.15]], [[0.5, 0.5], [0.8, 0.2],
+                                                [0.25, 0.75]]):
+        pure = stateless_random_ops.stateless_multinomial(
+            logits, num_samples, seed=seed_t)
+        values = [(seed, pure.eval(feed_dict={seed_t: seed})) for seed in seeds]
+        for s0, v0 in values:
+          for s1, v1 in values:
+            self.assertEqual(s0 == s1, np.all(v0 == v1))
+
+  def testEmpty(self):
+    with self.session():
+      with self.test_scope():
+        x = random_ops.multinomial(
+            array_ops.zeros([42, 40]), 0, output_dtype=dtypes.int32)
+        y = self.evaluate(x)
+        self.assertEqual(y.shape, (42, 0))
+
+  def testEmptyStateless(self):
+    with self.session() as sess:
+      with self.test_scope():
+        seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
+        x = stateless_random_ops.stateless_multinomial(
+            array_ops.zeros([42, 40]),
+            0,
+            seed=seed_t,
+            output_dtype=dtypes.int32)
+        y = sess.run(x, {seed_t: [0x12345678, 0xabcdef1]})
+        self.assertEqual(y.shape, (42, 0))
+
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/cholesky_op_test.py

@@ -0,0 +1,147 @@
+# Copyright 2015 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 tensorflow.ops.tf.Cholesky."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.platform import test
+
+
+class CholeskyOpTest(xla_test.XLATestCase):
+
+  # Cholesky defined for float64, float32, complex64, complex128
+  # (https://www.tensorflow.org/api_docs/python/tf/cholesky)
+  @property
+  def float_types(self):
+    return set(super(CholeskyOpTest, self).float_types).intersection(
+        (np.float64, np.float32, np.complex64, np.complex128))
+
+  def _verifyCholeskyBase(self, sess, placeholder, x, chol, verification, atol):
+    chol_np, verification_np = sess.run([chol, verification], {placeholder: x})
+    self.assertAllClose(x, verification_np, atol=atol)
+    self.assertShapeEqual(x, chol)
+    # Check that the cholesky is lower triangular, and has positive diagonal
+    # elements.
+    if chol_np.shape[-1] > 0:
+      chol_reshaped = np.reshape(chol_np, (-1, chol_np.shape[-2],
+                                           chol_np.shape[-1]))
+      for chol_matrix in chol_reshaped:
+        self.assertAllClose(chol_matrix, np.tril(chol_matrix), atol=atol)
+        self.assertTrue((np.diag(chol_matrix) > 0.0).all())
+
+  def _verifyCholesky(self, x, atol=1e-6):
+    # Verify that LL^T == x.
+    with self.session() as sess:
+      placeholder = array_ops.placeholder(
+          dtypes.as_dtype(x.dtype), shape=x.shape)
+      with self.test_scope():
+        chol = linalg_ops.cholesky(placeholder)
+      verification = test_util.matmul_without_tf32(chol, chol, adjoint_b=True)
+      self._verifyCholeskyBase(sess, placeholder, x, chol, verification, atol)
+
+  def testBasic(self):
+    data = np.array([[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]])
+    for dtype in self.float_types:
+      self._verifyCholesky(data.astype(dtype))
+
+  def testBatch(self):
+    for dtype in self.float_types:
+      simple_array = np.array(
+          [[[1., 0.], [0., 5.]]], dtype=dtype)  # shape (1, 2, 2)
+      self._verifyCholesky(simple_array)
+      self._verifyCholesky(np.vstack((simple_array, simple_array)))
+      odd_sized_array = np.array(
+          [[[4., -1., 2.], [-1., 6., 0], [2., 0., 5.]]], dtype=dtype)
+      self._verifyCholesky(np.vstack((odd_sized_array, odd_sized_array)))
+
+      # Generate random positive-definite matrices.
+      matrices = np.random.rand(10, 5, 5).astype(dtype)
+      for i in range(10):
+        matrices[i] = np.dot(matrices[i].T, matrices[i])
+      self._verifyCholesky(matrices, atol=1e-4)
+
+  @test_util.run_v2_only
+  def testNonSquareMatrixV2(self):
+    for dtype in self.float_types:
+      with self.assertRaises(errors.InvalidArgumentError):
+        linalg_ops.cholesky(np.array([[1., 2., 3.], [3., 4., 5.]], dtype=dtype))
+      with self.assertRaises(errors.InvalidArgumentError):
+        linalg_ops.cholesky(
+            np.array(
+                [[[1., 2., 3.], [3., 4., 5.]], [[1., 2., 3.], [3., 4., 5.]]],
+                dtype=dtype))
+
+  @test_util.run_v1_only("Different error types")
+  def testNonSquareMatrixV1(self):
+    for dtype in self.float_types:
+      with self.assertRaises(ValueError):
+        linalg_ops.cholesky(np.array([[1., 2., 3.], [3., 4., 5.]], dtype=dtype))
+      with self.assertRaises(ValueError):
+        linalg_ops.cholesky(
+            np.array(
+                [[[1., 2., 3.], [3., 4., 5.]], [[1., 2., 3.], [3., 4., 5.]]],
+                dtype=dtype))
+
+  @test_util.run_v2_only
+  def testWrongDimensionsV2(self):
+    for dtype in self.float_types:
+      tensor3 = constant_op.constant([1., 2.], dtype=dtype)
+      with self.assertRaises(errors.InvalidArgumentError):
+        linalg_ops.cholesky(tensor3)
+      with self.assertRaises(errors.InvalidArgumentError):
+        linalg_ops.cholesky(tensor3)
+
+  @test_util.run_v1_only("Different error types")
+  def testWrongDimensionsV1(self):
+    for dtype in self.float_types:
+      tensor3 = constant_op.constant([1., 2.], dtype=dtype)
+      with self.assertRaises(ValueError):
+        linalg_ops.cholesky(tensor3)
+      with self.assertRaises(ValueError):
+        linalg_ops.cholesky(tensor3)
+
+  def testLarge2000x2000(self):
+    n = 2000
+    shape = (n, n)
+    data = np.ones(shape).astype(np.float32) / (2.0 * n) + np.diag(
+        np.ones(n).astype(np.float32))
+    self._verifyCholesky(data, atol=1e-4)
+
+  def testMatrixConditionNumbers(self):
+    for dtype in self.float_types:
+      condition_number = 1000
+      size = 20
+
+      # Generate random positive-definite symmetric matrices, and take their
+      # Eigendecomposition.
+      matrix = np.random.rand(size, size)
+      matrix = np.dot(matrix.T, matrix)
+      _, w = np.linalg.eigh(matrix)
+
+      # Build new Eigenvalues exponentially distributed between 1 and
+      # 1/condition_number
+      v = np.exp(-np.log(condition_number) * np.linspace(0, size, size) / size)
+      matrix = np.dot(np.dot(w, np.diag(v)), w.T).astype(dtype)
+      self._verifyCholesky(matrix, atol=1e-4)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/clustering_test.py

@@ -0,0 +1,98 @@
+# 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.
+# ==============================================================================
+"""Tests for the behavior of the auto-compilation pass."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0"
+
+
+class ClusteringTest(xla_test.XLATestCase):
+
+  def testAdd(self):
+    val1 = np.array([4, 3, 2, 1], dtype=np.float32)
+    val2 = np.array([5, 6, 7, 8], dtype=np.float32)
+    expected = val1 + val2
+    with self.session():
+      with self.test_scope():
+        input1 = constant_op.constant(val1, name="const1")
+        input2 = constant_op.constant(val2, name="const2")
+        output = math_ops.add(input1, input2)
+      result = self.evaluate(output)
+    self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testAddFromCpuMultiple(self):
+    val1 = np.array([4, 3, 2, 1]).astype(np.float32)
+    val2 = np.array([5, 6, 7, 8]).astype(np.float32)
+    expected = val1 + val2
+    with self.session():
+      with ops.device(CPU_DEVICE):
+        input1 = constant_op.constant(val1, name="const1")
+        input2 = constant_op.constant(val2, name="const2")
+      with self.test_scope():
+        output = math_ops.add(input1, input2)
+      for _ in range(10):
+        result = self.evaluate(output)
+        self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testDeadlock(self):
+    # Builds a graph of the form:
+    #  x -> y
+    #       | \
+    #       z -> w
+    # where x and z are placed on the CPU and y and w are placed on the XLA
+    # device. If y and w are clustered for compilation, then the graph will
+    # deadlock since the clustered graph will contain a self-loop.
+    with self.session() as sess:
+      with ops.device(CPU_DEVICE):
+        x = array_ops.placeholder(dtypes.float32, [2])
+      with self.test_scope():
+        y = x * 2
+      with ops.device(CPU_DEVICE):
+        z = y * y
+      with self.test_scope():
+        w = y + z
+      result = sess.run(w, {x: [1.5, 0.5]})
+    self.assertAllClose(result, [12., 2.], rtol=1e-3)
+
+  def testHostMemory(self):
+    with self.session() as sess:
+      x = array_ops.placeholder(dtypes.int32)
+      with self.test_scope():
+        y = x + 1
+      with ops.device(CPU_DEVICE):
+        # Place a computation on the CPU, so y and w cannot be merged into the
+        # same JIT compilation.
+        z = y * 2
+      with self.test_scope():
+        # Argument 'y' is a non-constant output of a previous cluster. Make sure
+        # it is properly copied to host memory so it can be used as a
+        # compile-time constant input for this cluster.
+        w = array_ops.reshape(z, y)
+      result = sess.run(w, {x: [1, 0]})
+      expected = np.array([[4], [2]], dtype=np.int32)
+      self.assertAllClose(expected, result, rtol=1e-3)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/complex_div_test.py

@@ -0,0 +1,176 @@
+# 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.
+# ==============================================================================
+"""Test cases for complex numbers division."""
+
+import os
+
+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 gen_math_ops
+from tensorflow.python.platform import googletest
+
+os.environ["XLA_FLAGS"] = ("--xla_cpu_fast_math_honor_nans=true "
+                           "--xla_cpu_fast_math_honor_infs=true")
+
+
+class ComplexNumbersDivisionTest(xla_test.XLATestCase):
+  """Test cases for complex numbers division operators."""
+
+  def _testBinary(self, op, a, b, expected, equality_test=None):
+    with self.session() as session:
+      with self.test_scope():
+        pa = array_ops.placeholder(dtypes.as_dtype(a.dtype), a.shape, name="a")
+        pb = array_ops.placeholder(dtypes.as_dtype(b.dtype), b.shape, name="b")
+        output = op(pa, pb)
+      result = session.run(output, {pa: a, pb: b})
+      if equality_test is None:
+        equality_test = self.assertAllCloseAccordingToType
+      equality_test(np.real(result), np.real(expected), rtol=1e-3)
+      equality_test(np.imag(result), np.imag(expected), rtol=1e-3)
+
+  def testComplexOps(self):
+    for dtype in self.complex_types:
+      # Test division by 0 scenarios.
+      self._testBinary(
+          gen_math_ops.real_div,
+          np.array([
+              complex(1, 1),
+              complex(1, np.inf),
+              complex(1, np.nan),
+              complex(np.inf, 1),
+              complex(np.inf, np.inf),
+              complex(np.inf, np.nan),
+              complex(np.nan, 1),
+              complex(np.nan, np.inf),
+              complex(np.nan, np.nan),
+              complex(-np.inf, np.nan),
+          ],
+                   dtype=dtype),
+          np.array([
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0 + 0j,
+              0.0 + 0j,
+          ],
+                   dtype=dtype),
+          expected=np.array([
+              complex(np.inf, np.inf),
+              complex(np.inf, np.inf),
+              complex(np.inf, np.nan),
+              complex(np.inf, np.inf),
+              complex(np.inf, np.inf),
+              complex(np.inf, np.nan),
+              complex(np.nan, np.inf),
+              complex(np.nan, np.inf),
+              complex(np.nan, np.nan),
+              complex(-np.inf, np.nan),
+          ],
+                            dtype=dtype))
+
+      # Test division with finite numerator, inf/nan denominator.
+      self._testBinary(
+          gen_math_ops.real_div,
+          np.array([
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+          ],
+                   dtype=dtype),
+          np.array(
+              [
+                  complex(1, np.inf),
+                  complex(1, np.nan),
+                  complex(np.inf, 1),
+                  complex(np.inf, np.inf),  # C++ and Python diverge here.
+                  complex(np.inf, np.nan),  # C++ and Python diverge here.
+                  complex(np.nan, 1),
+                  complex(np.nan, np.inf),  # C++ and Python diverge here.
+                  complex(np.nan, -np.inf),  # C++ and Python diverge here.
+                  complex(np.nan, np.nan),
+              ],
+              dtype=dtype),
+          expected=np.array(
+              [
+                  (1 + 1j) / complex(1, np.inf),
+                  (1 + 1j) / complex(1, np.nan),
+                  (1 + 1j) / complex(np.inf, 1),
+                  complex(0 + 0j),  # C++ and Python diverge here.
+                  complex(0 + 0j),  # C++ and Python diverge here.
+                  (1 + 1j) / complex(np.nan, 1),
+                  complex(0 + 0j),  # C++ and Python diverge here.
+                  complex(0 - 0j),  # C++ and Python diverge here.
+                  (1 + 1j) / complex(np.nan, np.nan),
+              ],
+              dtype=dtype))
+
+      # Test division with inf/nan numerator, infinite denominator.
+      self._testBinary(
+          gen_math_ops.real_div,
+          np.array([
+              complex(1, np.inf),
+              complex(1, np.nan),
+              complex(np.inf, 1),
+              complex(np.inf, np.inf),
+              complex(np.inf, np.nan),
+              complex(np.nan, 1),
+              complex(np.nan, np.inf),
+              complex(np.nan, np.nan),
+              complex(np.nan, -np.inf),
+          ],
+                   dtype=dtype),
+          np.array([
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              1 + 1j,
+              -1 - 1j,
+          ],
+                   dtype=dtype),
+          expected=np.array(
+              [
+                  complex(np.inf, np.inf),  # C++ and Python diverge here.
+                  complex(1 / np.nan) / (1 + 1j),
+                  complex(np.inf / 1) / (1 + 1j),
+                  complex(np.inf, -np.nan),  # C++ and Python diverge here.
+                  complex(np.inf, -np.inf),  # C++ and Python diverge here.
+                  complex(np.nan / 1) / (1 + 1j),
+                  complex(np.inf, np.inf),  # C++ and Python diverge here.
+                  complex(np.nan / np.nan) / (1 + 1j),
+                  complex(np.inf, np.inf),  # C++ and Python diverge here.
+              ],
+              dtype=dtype))
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/concat_ops_test.py

@@ -0,0 +1,397 @@
+# 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 Concat Op."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python import pywrap_sanitizers
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class ConcatTest(xla_test.XLATestCase):
+
+  def testHStack(self):
+    with self.session():
+      p1 = array_ops.placeholder(dtypes.float32, shape=[4, 4])
+      p2 = array_ops.placeholder(dtypes.float32, shape=[4, 4])
+      with self.test_scope():
+        c = array_ops.concat([p1, p2], 0)
+      params = {
+          p1: np.random.rand(4, 4).astype("f"),
+          p2: np.random.rand(4, 4).astype("f")
+      }
+      result = c.eval(feed_dict=params)
+
+    self.assertEqual(result.shape, c.get_shape())
+    self.assertAllEqual(result[:4, :], params[p1])
+    self.assertAllEqual(result[4:, :], params[p2])
+
+  def testVStack(self):
+    with self.session():
+      p1 = array_ops.placeholder(dtypes.float32, shape=[4, 4])
+      p2 = array_ops.placeholder(dtypes.float32, shape=[4, 4])
+      with self.test_scope():
+        c = array_ops.concat([p1, p2], 1)
+      params = {
+          p1: np.random.rand(4, 4).astype("f"),
+          p2: np.random.rand(4, 4).astype("f")
+      }
+      result = c.eval(feed_dict=params)
+
+    self.assertEqual(result.shape, c.get_shape())
+    self.assertAllEqual(result[:, :4], params[p1])
+    self.assertAllEqual(result[:, 4:], params[p2])
+
+  def testInt32(self):
+    with self.session():
+      p1 = np.random.rand(2, 3).astype("i")
+      p2 = np.random.rand(2, 3).astype("i")
+      x1 = constant_op.constant(p1)
+      x2 = constant_op.constant(p2)
+      with self.test_scope():
+        c = array_ops.concat([x1, x2], 0)
+      result = self.evaluate(c)
+    self.assertAllEqual(result[:2, :], p1)
+    self.assertAllEqual(result[2:, :], p2)
+
+  def testAxisInt64(self):
+    with self.session():
+      p1 = np.random.rand(2, 3).astype("i")
+      p2 = np.random.rand(2, 3).astype("i")
+      x1 = constant_op.constant(p1)
+      x2 = constant_op.constant(p2)
+      axis = constant_op.constant(0, dtype=dtypes.int64)
+      with self.test_scope():
+        c = array_ops.concat([x1, x2], axis)
+      result = self.evaluate(c)
+    self.assertAllEqual(result[:2, :], p1)
+    self.assertAllEqual(result[2:, :], p2)
+
+  def _testRandom(self, dtype):
+    # Random dims of rank 5
+    shape = np.random.randint(1, 5, size=5)
+    # Random number of tensors, but always > 1.
+    num_tensors = np.random.randint(2, 10)
+    # Random dim to concat on
+    concat_dim = np.random.randint(5)
+    params = {}
+    if dtype == dtypes.bfloat16:
+      dtype_feed = dtypes.float32
+    else:
+      dtype_feed = dtype
+    with self.session():
+      p = []
+      for i in np.arange(num_tensors):
+        input_shape = shape
+        input_shape[concat_dim] = np.random.randint(1, 5)
+        placeholder = array_ops.placeholder(dtype_feed, shape=input_shape)
+        p.append(placeholder)
+
+        t = dtype_feed.as_numpy_dtype
+        params[placeholder] = np.random.rand(*input_shape).astype(t)
+
+      if dtype != dtype_feed:
+        concat_inputs = [math_ops.cast(p_i, dtype) for p_i in p]
+      else:
+        concat_inputs = p
+      with self.test_scope():
+        c = array_ops.concat(concat_inputs, concat_dim)
+        if dtype != dtype_feed:
+          c = math_ops.cast(c, dtype_feed)
+      result = c.eval(feed_dict=params)
+
+    self.assertEqual(result.shape, c.get_shape())
+    cur_offset = 0
+
+    for i in np.arange(num_tensors):
+      # The index into the result is the ':' along all dimensions
+      # except the concat_dim. slice(0, size) is used for ':', and
+      # a list of slices is used to index into result.
+      ind = [slice(0, params[p[i]].shape[j]) for j in np.arange(5)]
+      ind[concat_dim] = slice(cur_offset,
+                              cur_offset + params[p[i]].shape[concat_dim])
+      cur_offset += params[p[i]].shape[concat_dim]
+      if dtype == dtype_feed:
+        self.assertAllEqual(result[tuple(ind)], params[p[i]])
+      else:
+        self.assertAllClose(result[tuple(ind)], params[p[i]], 0.01)
+
+  def testRandom(self):
+    self._testRandom(dtypes.float32)
+    self._testRandom(dtypes.int32)
+
+  def _testGradientsSimple(self):
+    with self.session():
+      inp = []
+      inp_tensors = []
+      with self.test_scope():
+        for x in [1, 2, 6]:
+          shape = [10, x, 2]
+          t = np.random.rand(*shape).astype("f")
+          inp.append(t)
+          inp_tensors.append(
+              constant_op.constant(
+                  [float(y) for y in t.flatten()],
+                  shape=shape,
+                  dtype=dtypes.float32))
+        c = array_ops.concat(inp_tensors, 1)
+        output_shape = [10, 9, 2]
+        grad_inp = np.random.rand(*output_shape).astype("f")
+        grad_tensor = constant_op.constant(
+            [float(x) for x in grad_inp.flatten()], shape=output_shape)
+        grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor])
+        concated_grad = array_ops.concat(grad, 1)
+      result = self.evaluate(concated_grad)
+    self.assertAllEqual(result, grad_inp)
+
+  def testGradientsSimpleAll(self):
+    self._testGradientsSimple()
+
+  def _testGradientsFirstDim(self):
+    with self.session():
+      inp = []
+      inp_tensors = []
+      with self.test_scope():
+        for x in [1, 2, 6]:
+          shape = [x, 10, 2]
+          t = np.random.rand(*shape).astype("f")
+          inp.append(t)
+          inp_tensors.append(
+              constant_op.constant(
+                  [float(y) for y in t.flatten()],
+                  shape=shape,
+                  dtype=dtypes.float32))
+        c = array_ops.concat(inp_tensors, 0)
+        output_shape = [9, 10, 2]
+        grad_inp = np.random.rand(*output_shape).astype("f")
+        grad_tensor = constant_op.constant(
+            [float(x) for x in grad_inp.flatten()], shape=output_shape)
+        grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor])
+        concated_grad = array_ops.concat(grad, 0)
+        result = self.evaluate(concated_grad)
+
+    self.assertAllEqual(result, grad_inp)
+
+  def testGradientsFirstDimAll(self):
+    self._testGradientsFirstDim()
+
+  def _testGradientsLastDim(self):
+    with self.session():
+      inp = []
+      inp_tensors = []
+      with self.test_scope():
+        for x in [1, 2, 6]:
+          shape = [10, 2, x]
+          t = np.random.rand(*shape).astype("f")
+          inp.append(t)
+          inp_tensors.append(
+              constant_op.constant(
+                  [float(y) for y in t.flatten()],
+                  shape=shape,
+                  dtype=dtypes.float32))
+        c = array_ops.concat(inp_tensors, 2)
+        output_shape = [10, 2, 9]
+        grad_inp = np.random.rand(*output_shape).astype("f")
+        grad_tensor = constant_op.constant(
+            [float(x) for x in grad_inp.flatten()], shape=output_shape)
+        grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor])
+        concated_grad = array_ops.concat(grad, 2)
+        result = self.evaluate(concated_grad)
+
+    self.assertAllEqual(result, grad_inp)
+
+  def testGradientsLastDimAll(self):
+    self._testGradientsLastDim()
+
+  def _RunAndVerifyGradientsRandom(self):
+    # Random dims of rank 5
+    input_shape = np.random.randint(1, 5, size=5)
+    # Random number of tensors
+    num_tensors = np.random.randint(1, 10)
+    # Random dim to concat on
+    concat_dim = np.random.randint(5)
+    concat_dim_sizes = np.random.randint(1, 5, size=num_tensors)
+    with self.session():
+      inp = []
+      inp_tensors = []
+      with self.test_scope():
+        for x in concat_dim_sizes:
+          shape = input_shape
+          shape[concat_dim] = x
+          t = np.random.rand(*shape).astype("f")
+          inp.append(t)
+          inp_tensors.append(
+              constant_op.constant(
+                  [float(y) for y in t.flatten()],
+                  shape=shape,
+                  dtype=dtypes.float32))
+        c = array_ops.concat(inp_tensors, concat_dim)
+        output_shape = input_shape
+        output_shape[concat_dim] = concat_dim_sizes.sum()
+        grad_inp = np.random.rand(*output_shape).astype("f")
+        grad_tensor = constant_op.constant(
+            [float(x) for x in grad_inp.flatten()], shape=output_shape)
+        grad = gradients_impl.gradients([c], inp_tensors, [grad_tensor])
+        concated_grad = array_ops.concat(grad, concat_dim)
+        result = self.evaluate(concated_grad)
+
+    self.assertAllEqual(result, grad_inp)
+
+  def testGradientsRandom(self):
+    for _ in range(5):
+      self._RunAndVerifyGradientsRandom()
+
+  # Re-enable once zero-element Retvals are handled correctly.
+  def DISABLED_testZeroSize(self):
+    # Verify that concat doesn't crash and burn for zero size inputs
+    np.random.seed(7)
+    with self.session():
+      with self.test_scope():
+        for shape0 in (), (2,):
+          axis = len(shape0)
+          for shape1 in (), (3,):
+            for n0 in 0, 1, 2:
+              for n1 in 0, 1, 2:
+                x0 = np.random.randn(*(shape0 + (n0,) + shape1))
+                x1 = np.random.randn(*(shape0 + (n1,) + shape1))
+                correct = np.concatenate([x0, x1], axis=axis)
+                # TODO(irving): Make tf.concat handle map, then drop list().
+                xs = list(map(constant_op.constant, [x0, x1]))
+                c = array_ops.concat(xs, axis)
+                self.assertAllEqual(c, correct)
+                # Check gradients
+                dc = np.random.randn(*c.get_shape().as_list())
+                dxs = self.evaluate(gradients_impl.gradients(c, xs, dc))
+                self.assertAllEqual(dc, np.concatenate(dxs, axis=axis))
+
+  def testConcatTuple(self):
+    c1 = np.random.rand(4, 4).astype(np.float32)
+    c2 = np.random.rand(4, 4).astype(np.float32)
+    with self.session():
+      with self.test_scope():
+        concat_list_t = array_ops.concat([c1, c2], 0)
+        concat_tuple_t = array_ops.concat((c1, c2), 0)
+      self.assertAllEqual(concat_list_t, self.evaluate(concat_tuple_t))
+
+  def testConcatNoScalars(self):
+    with self.session():
+      with self.test_scope():
+        scalar = constant_op.constant(7)
+        dim = array_ops.placeholder(dtypes.int32)
+        with self.assertRaisesRegex(
+            ValueError, r"Can't concatenate scalars \(use tf\.stack instead\)"):
+          array_ops.concat([scalar, scalar, scalar], dim)
+
+  # The purpose of this is to ensure that XLA on GPU will not run out of memory
+  # with too many arguments.
+  def testConcatLargeNumberOfTensors(self):
+    if "CPU" in self.device:
+      self.skipTest("This test can time out on CPU, so we will just allow "
+                    "other backends to catch this specific error.")
+    if (pywrap_sanitizers.is_asan_enabled() or
+        pywrap_sanitizers.is_tsan_enabled() or
+        pywrap_sanitizers.is_msan_enabled() or
+        pywrap_sanitizers.is_ubsan_enabled()):
+      self.skipTest("This test can time out on *SAN.")
+    with self.session():
+      with self.test_scope():
+        for concat_dim in range(2):
+          params = {}
+          p = []
+          shape = np.array([7, 13])
+          num_tensors = 1001
+          for i in np.arange(num_tensors):
+            input_shape = shape
+            placeholder = array_ops.placeholder(
+                dtypes.float32, shape=input_shape)
+            p.append(placeholder)
+            params[placeholder] = np.random.rand(*input_shape).astype(
+                np.float32)
+
+          concat_inputs = p
+          c = array_ops.concat(concat_inputs, concat_dim)
+          result = c.eval(feed_dict=params)
+
+          self.assertEqual(result.shape, c.get_shape())
+          cur_offset = 0
+
+          for i in np.arange(num_tensors):
+            # The index into the result is the ':' along all dimensions
+            # except the concat_dim. slice(0, size) is used for ':', and
+            # a list of slices is used to index into result.
+            index = [slice(0, params[p[i]].shape[j]) for j in np.arange(2)]
+            index[concat_dim] = slice(
+                cur_offset, cur_offset + params[p[i]].shape[concat_dim])
+            cur_offset += params[p[i]].shape[concat_dim]
+            self.assertAllEqual(result[tuple(index)], params[p[i]])
+
+
+class ConcatOffsetTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    with self.session():
+      with self.test_scope():
+        cdim = constant_op.constant(1, dtypes.int32)
+        s0 = constant_op.constant([2, 3, 5], dtypes.int32)
+        s1 = constant_op.constant([2, 7, 5], dtypes.int32)
+        s2 = constant_op.constant([2, 20, 5], dtypes.int32)
+        off = gen_array_ops.concat_offset(cdim, [s0, s1, s2])
+        ans = self.evaluate(off)
+        self.assertAllEqual(ans, [[0, 0, 0], [0, 3, 0], [0, 10, 0]])
+
+
+class PackTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    with self.session():
+      with self.test_scope():
+        s0 = constant_op.constant([2, 3, 5], dtypes.int32)
+        s1 = constant_op.constant([2, 7, 5], dtypes.int32)
+        s2 = constant_op.constant([2, 20, 5], dtypes.int32)
+        packed = array_ops_stack.stack([s0, s1, s2])
+        ans = self.evaluate(packed)
+        self.assertAllEqual(ans, [[2, 3, 5], [2, 7, 5], [2, 20, 5]])
+
+  def testScalars(self):
+    with self.session():
+      with self.test_scope():
+        s0 = constant_op.constant(2, dtypes.int32)
+        s1 = constant_op.constant(3, dtypes.int32)
+        s2 = constant_op.constant(5, dtypes.int32)
+        packed = array_ops_stack.stack([s0, s1, s2])
+        ans = self.evaluate(packed)
+        self.assertAllEqual(ans, [2, 3, 5])
+
+  def testEmpty(self):
+    with self.session():
+      with self.test_scope():
+        s0 = constant_op.constant([[]], dtypes.int32)
+        s1 = constant_op.constant([[]], dtypes.int32)
+        s2 = constant_op.constant([[]], dtypes.int32)
+        packed = array_ops_stack.stack([s0, s1, s2])
+        ans = self.evaluate(packed)
+        self.assertAllEqual(ans, [[[]], [[]], [[]]])
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/cond_test.py

@@ -0,0 +1,318 @@
+# 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 tf.cond in XLA."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.client import session
+from tensorflow.python.compiler.xla import xla
+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 errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_switch_case
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.platform import test
+
+
+@test_util.with_control_flow_v2
+class CondTest(xla_test.XLATestCase):
+
+  def testCondAndTensorArrayInDefun(self):
+    # TODO(b/132430685): Make test more useful. Also b/129396295, b/127846988
+    with self.session(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      @def_function.function
+      def f():
+        ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=1)
+        output = cond.cond(
+            constant_op.constant(True),
+            lambda: ta.write(0, 5.), lambda: ta.write(0, 10.))
+
+        return output.stack()
+
+      output_t = f()
+      self.assertAllEqual([5.], self.evaluate(output_t))
+
+      xla_context.Exit()
+
+  def testCondAndTensorArrayInDefun_constFolding(self):
+    g = ops.Graph()
+    with session.Session(graph=g), g.as_default(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      @def_function.function
+      def f():
+        ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=1)
+        output = cond.cond(
+            constant_op.constant(False),
+            lambda: ta.write(0, 5.), lambda: ta.write(0, 10.))
+
+        return output.stack()
+
+      output_t = f()
+      self.assertAllEqual([10.], self.evaluate(output_t))
+
+      xla_context.Exit()
+
+  def testCondAndTensorArray_xlaCompile(self):
+    self.skipTest("b/127846988")
+    # Fails with "Uninitialized arguments" in XlaIfOp::Compile
+    with self.session(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      def f():
+        ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=1)
+        output = cond.cond(
+            constant_op.constant(True),
+            lambda: ta.write(0, 5.), lambda: ta.write(0, 10.))
+
+        return output.stack()
+
+      output_t, = xla.compile(f)
+      self.assertAllEqual([5.], self.evaluate(output_t))
+
+      xla_context.Exit()
+
+  def testCondConstPropagation(self):
+    with self.session() as sess, self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      x = array_ops.placeholder(dtypes.float32)
+      p = array_ops.placeholder(dtypes.int32)
+
+      # TODO(b/129021699): Wrapping this in a tf.function does not work.
+      def if_true():
+        # This emits a StridedSlice op which expects the index to be a
+        # compile-time const.
+        return x[p]
+
+      def if_false():
+        return 5.
+
+      output = cond.cond(
+          constant_op.constant(True), if_true, if_false)
+
+      self.assertAllEqual(1.,
+                          sess.run(output, feed_dict={
+                              x: [0., 1., 2.],
+                              p: 1
+                          }))
+
+      xla_context.Exit()
+
+  def testCondConstPropagation_xlaCompile(self):
+    self.skipTest("b/132430685")
+    with self.session(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      x = array_ops.placeholder_with_default([0., 1., 2.], shape=[3])
+      p = constant_op.constant(1)
+
+      def f():
+        # TODO(b/129021699): Wrapping this in a tf.function does not work.
+        def if_true():
+          # This emits a StridedSlice op which expects the index to be a
+          # compile-time const.
+          return x[p]
+
+        def if_false():
+          return 5.
+
+        return cond.cond(
+            constant_op.constant(True), if_true, if_false)
+
+      output = xla.compile(f)
+
+      self.assertAllEqual(1., self.evaluate(output))
+
+      xla_context.Exit()
+
+  def testCondConstPropagation_errorMsg(self):
+    self.skipTest("b/132430685")
+    with self.session() as sess, self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      x = array_ops.placeholder(dtypes.float32)
+      p = random_ops.random_uniform([], minval=1, maxval=3, dtype=dtypes.int32)
+
+      # TODO(b/129021699): Wrapping this in a tf.function does not work.
+      def if_true():
+        # This emits a StridedSlice op which expects the index to be a
+        # compile-time const.
+        return x[:p]
+
+      def if_false():
+        return array_ops.fill([p], 5.)
+
+      output = cond.cond(
+          constant_op.constant(True), if_true, if_false)
+
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  "must be a compile-time constant"):
+        sess.run(
+            output, feed_dict={
+                x: [0., 1., 2.],
+            })
+
+      xla_context.Exit()
+
+  def testCondConstPropagation_errorMsg_xlaCompile(self):
+    with self.session() as sess, self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      x = array_ops.placeholder(dtypes.float32)
+      p = random_ops.random_uniform([], minval=1, maxval=3, dtype=dtypes.int32)
+      condition = math_ops.cast(
+          random_ops.random_uniform([], minval=0, maxval=2, dtype=dtypes.int32),
+          dtypes.bool)
+
+      def f():
+        # TODO(b/129021699): Wrapping this in a tf.function does not work.
+        def if_true():
+          # This emits a StridedSlice op which expects the index to be a
+          # compile-time const.
+          return x[:p]
+
+        def if_false():
+          return array_ops.fill([p], 5.)
+
+        return cond.cond(condition, if_true, if_false)
+
+      output = xla.compile(f)
+
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  "must be a compile-time constant"):
+        sess.run(
+            output, feed_dict={
+                x: [0., 1., 2.],
+            })
+
+      xla_context.Exit()
+
+  def testSwitchCaseAndTensorArrayInDefun(self):
+    self.skipTest("b/127846988")
+    with self.session(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      @def_function.function
+      def f():
+        ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=1)
+        output = control_flow_switch_case.switch_case(
+            constant_op.constant(1), {
+                0: lambda: ta.write(0, 5.),
+                1: lambda: ta.write(0, 10.),
+                2: lambda: ta.write(0, 15.),
+            })
+
+        return output.stack()
+
+      output_t = f()
+      self.assertAllEqual([10.], self.evaluate(output_t))
+
+      xla_context.Exit()
+
+  def testSwitchCaseAndTensorArray_xlaCompile(self):
+    self.skipTest("b/127846988")
+    with self.session(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      def f():
+        ta = tensor_array_ops.TensorArray(dtype=dtypes.float32, size=1)
+        output = control_flow_switch_case.switch_case(
+            constant_op.constant(1), {
+                0: lambda: ta.write(0, 5.),
+                1: lambda: ta.write(0, 10.),
+                2: lambda: ta.write(0, 15.),
+            })
+
+        return output.stack()
+
+      output_t, = xla.compile(f)
+      self.assertAllEqual([10.], self.evaluate(output_t))
+
+      xla_context.Exit()
+
+  def testSwitchCaseConstPropagation(self):
+    self.skipTest("b/127846988")
+    with self.session() as sess, self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      x = array_ops.placeholder(dtypes.float32)
+      p = array_ops.placeholder(dtypes.int32)
+
+      def branch0():
+        return 5.
+
+      def branch1():
+        return 15.
+
+      # TODO(b/129021699): Wrapping this in a tf.function does not work.
+      def branch2():
+        # This emits a StridedSlice op which expects the index to be a
+        # compile-time const.
+        return x[p]
+
+      output = control_flow_switch_case.switch_case(
+          constant_op.constant(2), {
+              0: branch0,
+              1: branch1,
+              2: branch2,
+          })
+
+      self.assertAllEqual(7.,
+                          sess.run(output, feed_dict={
+                              x: [0., 1., 7.],
+                              p: 2,
+                          }))
+
+      xla_context.Exit()
+
+  def testCondNoInputs(self):
+    """Verifies against `Failed precondition: Expected one input shape`."""
+
+    with self.session(), self.test_scope():
+      xla_context = control_flow_ops.XLAControlFlowContext()
+      xla_context.Enter()
+
+      for pred in True, False:
+        cond_out = cond.cond(
+            array_ops.placeholder_with_default(pred, []),
+            lambda: constant_op.constant(2.),
+            lambda: constant_op.constant(1.))
+        self.assertEqual(int(pred) + 1., self.evaluate(cond_out))
+
+      xla_context.Exit()
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/tests/const_arg_test.py

@@ -0,0 +1,42 @@
+# Copyright 2021 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 compilation that involves constant arguments."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.compiler.tf2xla.python import xla
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import googletest
+
+
+class ConstArgTest(xla_test.XLATestCase):
+  # Pass constant value to XlaDynamicSlice's size parameter that must be found
+  # by xla::ValueInference. Most often, constants are passed to op kernels
+  # using XlaExpression with kind kConstant. To require value inference, this
+  # model obfuscates the constant using operations `>=` and `where_v2`.
+
+  def testValueInference(self):
+    with self.session() as session:
+      with self.test_scope():
+        a = array_ops.placeholder(dtypes.int32, [], name="a")
+        size = array_ops.reshape(array_ops.where_v2(a >= 0, 1, 0), [1])
+        output = xla.dynamic_slice([11, 12, 13], [0], size)
+      result = session.run(output, {a: 1})
+      expected = [11]
+      self.assertEqual(result, expected)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/const_test.py

@@ -0,0 +1,81 @@
+# Copyright 2022 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 const op compilation."""
+
+import numpy as np
+
+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.platform import test
+
+
+# This test doesn't use XLATestCase like the other tests in this directory.
+# The Const op xla op kernel is compilation only and therefore is not executed
+# with XLA in the on demand compilation mode. Instead we use
+# tf.function(jit_compile=True)
+class ConstOpTest(test_util.TensorFlowTestCase):
+
+  # Verifies that the Const op works
+  # @test_util.run_v2_only
+  def testConst(self):
+    types = {
+        dtypes.bool,
+        dtypes.int8,
+        dtypes.int16,
+        dtypes.int32,
+        dtypes.int64,
+        dtypes.uint8,
+        dtypes.uint16,
+        dtypes.uint32,
+        dtypes.uint64,
+        dtypes.float16,
+        dtypes.bfloat16,
+        dtypes.float32,
+        dtypes.float64,
+        dtypes.float8_e5m2,
+        dtypes.float8_e4m3fn,
+        dtypes.float8_e4m3fnuz,
+        dtypes.float8_e4m3b11fnuz,
+        dtypes.float8_e5m2fnuz,
+    }
+    for dtype in types:
+      with self.subTest(dtype=dtype):
+        if dtype == dtypes.bool:
+          values = [True, False]
+        elif dtype in [
+            dtypes.uint8,
+            dtypes.uint16,
+            dtypes.uint32,
+            dtypes.uint64,
+        ]:
+          values = [0., 1., dtype.min, dtype.max]
+        else:
+          values = [0., 1., -1., dtype.min, dtype.max]
+        if dtype.is_floating:
+          values.extend([float("Inf"), -float("Inf"), float("NaN")])
+        values = np.array(values, dtype=dtype.as_numpy_dtype)
+
+        @def_function.function(jit_compile=True)
+        def f():
+          return constant_op.constant(values, dtype)  # pylint: disable=cell-var-from-loop
+
+        result = f()
+        self.assertAllEqual(self.evaluate(result), values)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/conv2d_test.py

@@ -0,0 +1,963 @@
+# 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.
+# ==============================================================================
+"""Tests for Conv2D via the XLA JIT.
+
+The canned results in these tests are created by running each test using the
+Tensorflow CPU device and saving the output.
+"""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import test_utils
+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 gen_nn_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import googletest
+
+DATA_FORMATS = (
+    ("_data_format_NHWC", "NHWC"),
+    ("_data_format_NCHW", "NCHW"),
+)
+
+CONV_CONFIGS = (
+    ("_Conv2D_data_format_NHWC", "NHWC", "Conv2D"),
+    ("_Conv2D_data_format_NCHW", "NCHW", "Conv2D"),
+    ("_Conv_data_format_NHWC", "NHWC", "Conv"),
+    ("_Conv_data_format_NCHW", "NCHW", "Conv"),
+)
+
+
+class Conv2DTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _VerifyValues(
+      self,
+      input_sizes=None,
+      filter_sizes=None,
+      strides=None,
+      dilations=None,
+      padding=None,
+      data_format_src="NHWC",
+      data_format_dst="NHWC",
+      expected=None,
+      op_name="Conv2D",
+  ):
+    """Tests that tf.nn.conv2d produces the expected value.
+
+    Args:
+      input_sizes: Input tensor dimensions in [batch, input_rows, input_cols,
+        input_depth].
+      filter_sizes: Filter tensor dimensions in [kernel_rows, kernel_cols,
+        input_depth, output_depth].
+      strides: Strides.
+      dilations: RHS dilations.
+      padding: Padding type.
+      data_format_src: Data format input is in.
+      data_format_dst: Data format verification will run and input is converted
+        to.
+      expected: Expected output.
+      op_name: Name of operation to test (Conv/Conv2D)
+    """
+
+    total_size_1 = np.prod(input_sizes)
+    total_size_2 = np.prod(filter_sizes)
+    x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(input_sizes)
+    x2 = np.arange(1, total_size_2 + 1, dtype=np.float32).reshape(filter_sizes)
+    strides = [1] + strides + [1]
+    if dilations is None:
+      dilations = [1, 1]
+    dilations = [1] + dilations + [1]
+
+    # Convert between data formats.
+    expected = test_utils.ConvertBetweenDataFormats(expected, data_format_src,
+                                                    data_format_dst)
+    x1 = test_utils.ConvertBetweenDataFormats(x1, data_format_src,
+                                              data_format_dst)
+    input_sizes = test_utils.PermuteDimsBetweenDataFormats(
+        input_sizes, data_format_src, data_format_dst)
+    strides = test_utils.PermuteDimsBetweenDataFormats(strides, data_format_src,
+                                                       data_format_dst)
+    dilations = test_utils.PermuteDimsBetweenDataFormats(
+        dilations, data_format_src, data_format_dst)
+
+    with self.session() as sess:
+      t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes)
+      t2 = array_ops.placeholder(dtypes.float32, shape=filter_sizes)
+      with self.test_scope():
+        if op_name == "Conv":
+          conv_format = (
+              "CHANNELS_LAST" if data_format_dst == "NHWC" else "CHANNELS_FIRST"
+          )
+          out = gen_nn_ops.conv(
+              t1,
+              t2,
+              strides=strides,
+              padding=padding,
+              data_format=conv_format,
+              dilations=dilations,
+          )
+        elif op_name == "Conv2D":
+          out = nn_ops.conv2d(
+              t1,
+              t2,
+              strides=strides,
+              padding=padding,
+              data_format=data_format_dst,
+              dilations=dilations,
+          )
+        else:
+          raise ValueError("Invalid op name: %s" % op_name)
+
+      value = sess.run(out, {t1: x1, t2: x2})
+      self.assertAllClose(expected, value, 1e-3)
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D1x1Filter(self, data_format, op_name):
+    expected_output = np.reshape([
+        30.0, 36.0, 42.0, 66.0, 81.0, 96.0, 102.0, 126.0, 150.0, 138.0, 171.0,
+        204.0, 174.0, 216.0, 258.0, 210.0, 261.0, 312.0
+    ], [1, 2, 3, 3])
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[1, 1, 3, 3],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D2x2Filter(self, data_format, op_name):
+    expected_output = np.reshape(
+        [2271.0, 2367.0, 2463.0, 2901.0, 3033.0, 3165.0], [1, 1, 2, 3])
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D2x2Filter2x1Dilation(self, data_format, op_name):
+    expected_output = np.array([[[[72], [82], [92]], [[112], [122], [132]]]])
+    self._VerifyValues(
+        input_sizes=[1, 4, 4, 1],
+        filter_sizes=[2, 2, 1, 1],
+        strides=[1, 1],
+        dilations=[2, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D1x2Filter(self, data_format, op_name):
+    expected_output = np.reshape([
+        231.0, 252.0, 273.0, 384.0, 423.0, 462.0, 690.0, 765.0, 840.0, 843.0,
+        936.0, 1029.0
+    ], [1, 2, 2, 3])
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[1, 2, 3, 3],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D2x2FilterStride2(self, data_format, op_name):
+    expected_output = np.reshape([2271.0, 2367.0, 2463.0], [1, 1, 1, 3])
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        strides=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D2x2FilterStride2Same(self, data_format, op_name):
+    expected_output = np.reshape(
+        [2271.0, 2367.0, 2463.0, 1230.0, 1305.0, 1380.0], [1, 1, 2, 3])
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        strides=[2, 2],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2DEmptyDilation(self, data_format, op_name):
+    self._VerifyValues(
+        input_sizes=[0, 2, 3, 3],
+        filter_sizes=[1, 1, 3, 3],
+        strides=[1, 1],
+        dilations=[2, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.zeros([0, 2, 3, 3]),
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D2x2FilterDilation(self, data_format, op_name):
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        strides=[1, 1],
+        dilations=[1, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.reshape([2667, 2781, 2895], [1, 1, 1, 3]),
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2D1x2FilterDilation(self, data_format, op_name):
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[1, 2, 3, 3],
+        strides=[1, 1],
+        dilations=[2, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.array([[
+            [[231, 252, 273], [384, 423, 462]],
+            [[690, 765, 840], [843, 936, 1029]],
+        ]]),
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv2DKernelSizeMatchesInputSizeDilation(self, data_format, op_name):
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 1],
+        filter_sizes=[2, 2, 1, 2],
+        strides=[1, 1],
+        dilations=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.reshape([108, 128], [1, 1, 1, 2]),
+        op_name=op_name,
+    )
+
+  def testConvExpandedBatch(self):
+    tensor_in_sizes_batch = [10, 2, 3, 3]
+    tensor_in_sizes_expanded_batch = [2, 5, 2, 3, 3]
+    batch_dims = 2
+    filter_in_sizes = [1, 1, 3, 3]
+    filter_in = np.arange(
+        1, np.prod(filter_in_sizes) + 1, dtype=np.float32
+    ).reshape(filter_in_sizes)
+    x1 = np.arange(
+        1, np.prod(tensor_in_sizes_batch) + 1, dtype=np.float32
+    ).reshape(tensor_in_sizes_batch)
+    x2 = x1.reshape(tensor_in_sizes_expanded_batch)
+
+    with self.session() as sess:
+      t1 = array_ops.placeholder(dtypes.bfloat16, shape=tensor_in_sizes_batch)
+      t2 = array_ops.placeholder(
+          dtypes.bfloat16, shape=tensor_in_sizes_expanded_batch
+      )
+      filter_t = array_ops.placeholder(dtypes.bfloat16, shape=filter_in_sizes)
+
+      out1 = gen_nn_ops.conv(
+          t1, filter_t, strides=[1, 1, 1, 1], padding="VALID"
+      )
+      out2 = gen_nn_ops.conv(
+          t2,
+          filter_t,
+          strides=[1, 1, 1, 1],
+          padding="VALID",
+          batch_dims=batch_dims,
+      )
+      value1 = sess.run(out1, {t1: x1, filter_t: filter_in})
+      value2 = sess.run(out2, {t2: x2, filter_t: filter_in})
+
+      self.assertEqual(list(value1.shape), tensor_in_sizes_batch)
+      self.assertEqual(list(value2.shape), tensor_in_sizes_expanded_batch)
+      self.assertAllCloseAccordingToType(value1, value2.reshape(value1.shape))
+
+
+class Conv2DBackpropInputTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _VerifyValues(self,
+                    input_sizes=None,
+                    filter_sizes=None,
+                    out_backprop_sizes=None,
+                    strides=None,
+                    dilations=None,
+                    padding=None,
+                    data_format_src="NHWC",
+                    data_format_dst="NHWC",
+                    expected=None):
+    """Tests that gen_nn_ops.conv2d_backprop_input produces the expected output.
+
+    Args:
+      input_sizes: Input tensor dimensions in
+        [batch, input_rows, input_cols, input_depth].
+      filter_sizes: Filter tensor dimensions in
+        [kernel_rows, kernel_cols, input_depth, output_depth].
+      out_backprop_sizes: Output gradients tensor dimensions.
+      strides: Strides.
+      dilations: Dilations.
+      padding: Padding type.
+      data_format_src: Data format input is in.
+      data_format_dst: Data format verification will run and input is converted
+        to.
+      expected: Expected output.
+    """
+
+    total_size_1 = np.prod(filter_sizes)
+    total_size_2 = np.prod(out_backprop_sizes)
+    x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(filter_sizes)
+    x2 = np.arange(
+        1, total_size_2 + 1, dtype=np.float32).reshape(out_backprop_sizes)
+    strides = [1] + strides + [1]
+    if dilations is not None:
+      dilations = [1] + dilations + [1]
+
+    expected = np.reshape(expected, input_sizes)
+
+    # Convert between data formats.
+    expected = test_utils.ConvertBetweenDataFormats(expected, data_format_src,
+                                                    data_format_dst)
+    x2 = test_utils.ConvertBetweenDataFormats(x2, data_format_src,
+                                              data_format_dst)
+    input_sizes = test_utils.PermuteDimsBetweenDataFormats(
+        input_sizes, data_format_src, data_format_dst)
+    out_backprop_sizes = test_utils.PermuteDimsBetweenDataFormats(
+        out_backprop_sizes, data_format_src, data_format_dst)
+    strides = test_utils.PermuteDimsBetweenDataFormats(strides, data_format_src,
+                                                       data_format_dst)
+    if dilations is not None:
+      dilations = test_utils.PermuteDimsBetweenDataFormats(
+          dilations, data_format_src, data_format_dst)
+
+    with self.session() as sess:
+      t1 = array_ops.placeholder(dtypes.float32, shape=filter_sizes)
+      t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes)
+      with self.test_scope():
+        out = gen_nn_ops.conv2d_backprop_input(
+            input_sizes=input_sizes,
+            filter=t1,
+            out_backprop=t2,
+            strides=strides,
+            dilations=dilations,
+            padding=padding,
+            data_format=data_format_dst)
+
+      value = sess.run(out, {t1: x1, t2: x2})
+      self.assertAllEqual(input_sizes, value.shape)
+      self.assertAllClose(expected, value, 1e-3)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x1Filter(self, data_format):
+    expected_output = [
+        5, 11, 17, 11, 25, 39, 17, 39, 61, 23, 53, 83, 29, 67, 105, 35, 81, 127,
+        41, 95, 149, 47, 109, 171, 53, 123, 193, 59, 137, 215, 65, 151, 237, 71,
+        165, 259, 77, 179, 281, 83, 193, 303, 89, 207, 325, 95, 221, 347.
+    ]
+    self._VerifyValues(
+        input_sizes=[1, 4, 4, 3],
+        filter_sizes=[1, 1, 3, 2],
+        out_backprop_sizes=[1, 4, 4, 2],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterStride3Width5(self, data_format):
+    expected_output = [1, 2, 0, 2, 4]
+    self._VerifyValues(
+        input_sizes=[1, 1, 5, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[3, 3],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterStride3Width6(self, data_format):
+    expected_output = [1, 2, 0, 2, 4, 0]
+    self._VerifyValues(
+        input_sizes=[1, 1, 6, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[3, 3],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterStride3Width7(self, data_format):
+    expected_output = [1, 2, 0, 2, 4, 0, 0]
+    self._VerifyValues(
+        input_sizes=[1, 1, 7, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[3, 3],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterC1Same(self, data_format):
+    expected_output = [1, 4, 7, 7, 23, 33]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 2, 3, 1],
+        strides=[1, 1],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Filter(self, data_format):
+    expected_output = [
+        14, 32, 50, 100, 163, 226, 167, 212, 257, 122, 140, 158, 478, 541, 604,
+        437, 482, 527
+    ]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        out_backprop_sizes=[1, 1, 2, 3],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterSame(self, data_format):
+    expected_output = [
+        14, 32, 50, 100, 163, 226, 217, 334, 451, 190, 307, 424, 929, 1217,
+        1505, 1487, 1883, 2279
+    ]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        out_backprop_sizes=[1, 2, 3, 3],
+        strides=[1, 1],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2Filter(self, data_format):
+    expected_output = [1, 4, 4, 3, 10, 8, 5, 16, 12]
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 3, 2, 1],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterSame(self, data_format):
+    expected_output = [1, 4, 7, 4, 13, 16, 7, 22, 25]
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 3, 3, 1],
+        strides=[1, 1],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterStride2(self, data_format):
+    expected_output = [1, 2, 5, 4, 6, 0, 0, 0, 0, 0, 3, 6, 13, 8, 12]
+    self._VerifyValues(
+        input_sizes=[1, 3, 5, 1],
+        filter_sizes=[1, 3, 1, 1],
+        out_backprop_sizes=[1, 2, 2, 1],
+        strides=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterStride2Same(self, data_format):
+    expected_output = [1, 2, 2, 3, 4, 6]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[2, 2],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Depth3ValidBackpropInputStride1x1Dilation2x1(
+      self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 3, 6, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 5, 1],
+        strides=[1, 1],
+        dilations=[2, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[1, 4, 7, 10, 13, 10, 0, 0, 0, 0, 0, 0, 3, 10, 17, 24, 31, 20])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Depth1ValidBackpropInputDilation1x2(self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 1, 1],
+        strides=[1, 1],
+        dilations=[1, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[1, 0, 2, 3, 0, 4])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DEmptyBackpropInputDilation1x2(self, data_format):
+    self._VerifyValues(
+        input_sizes=[0, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[0, 1, 1, 1],
+        strides=[1, 1],
+        dilations=[1, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.zeros([0]))
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Depth3ValidBackpropInputDilation2x1(self, data_format):
+    # The GPU version of this test is not very stable. So adjusting the
+    # error threshold to 1e-4.
+    self._VerifyValues(
+        input_sizes=[1, 3, 2, 3],
+        filter_sizes=[2, 2, 3, 3],
+        out_backprop_sizes=[1, 1, 1, 3],
+        strides=[1, 1],
+        dilations=[2, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[
+            14, 32, 50, 68, 86, 104, 0, 0, 0, 0, 0, 0, 122, 140, 158, 176, 194,
+            212
+        ])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DKernelSizeMatchesInputSizeBackpropInputDilation2x2(
+      self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 1],
+        filter_sizes=[2, 2, 1, 2],
+        out_backprop_sizes=[1, 1, 1, 2],
+        strides=[1, 1],
+        dilations=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[5, 0, 11, 0, 0, 0, 17, 0, 23])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DGroupedFilter(self, data_format):
+    expected_output = [
+        5, 17, 29, 25, 53, 81, 41, 53, 65, 109, 137, 165, 77, 89, 101, 193, 221,
+        249, 113, 125, 137, 277, 305, 333
+    ]
+    self._VerifyValues(
+        input_sizes=[1, 2, 2, 6],
+        filter_sizes=[2, 2, 3, 4],
+        out_backprop_sizes=[1, 1, 1, 4],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+
+class Conv2DBackpropFilterTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _VerifyValues(self,
+                    input_sizes=None,
+                    filter_sizes=None,
+                    out_backprop_sizes=None,
+                    strides=None,
+                    dilations=None,
+                    padding=None,
+                    data_format_src="NHWC",
+                    data_format_dst="NHWC",
+                    expected=None):
+    """Tests that gen_nn_ops.conv2d_backprop_filter produces the right output.
+
+    Args:
+      input_sizes: Input tensor dimensions in
+        [batch, input_rows, input_cols, input_depth].
+      filter_sizes: Filter tensor dimensions in
+        [kernel_rows, kernel_cols, input_depth, output_depth].
+      out_backprop_sizes: Output gradients tensor dimensions.
+      strides: Stride.
+      dilations: Dilations.
+      padding: Padding type.
+      data_format_src: Data format input is in.
+      data_format_dst: Data format verification will run and input is converted
+        to.
+      expected: Expected output.
+    """
+
+    total_size_1 = np.prod(input_sizes)
+    total_size_2 = np.prod(out_backprop_sizes)
+    x1 = np.arange(1, total_size_1 + 1, dtype=np.float32).reshape(input_sizes)
+    x2 = np.arange(
+        1, total_size_2 + 1, dtype=np.float32).reshape(out_backprop_sizes)
+    strides = [1] + strides + [1]
+    if dilations is not None:
+      dilations = [1] + dilations + [1]
+
+    expected = np.reshape(expected, filter_sizes)
+
+    # Convert between data formats.
+    x1 = test_utils.ConvertBetweenDataFormats(x1, data_format_src,
+                                              data_format_dst)
+    x2 = test_utils.ConvertBetweenDataFormats(x2, data_format_src,
+                                              data_format_dst)
+    input_sizes = test_utils.PermuteDimsBetweenDataFormats(
+        input_sizes, data_format_src, data_format_dst)
+    out_backprop_sizes = test_utils.PermuteDimsBetweenDataFormats(
+        out_backprop_sizes, data_format_src, data_format_dst)
+    strides = test_utils.PermuteDimsBetweenDataFormats(strides, data_format_src,
+                                                       data_format_dst)
+    if dilations is not None:
+      dilations = test_utils.PermuteDimsBetweenDataFormats(
+          dilations, data_format_src, data_format_dst)
+
+    with self.session() as sess:
+      t1 = array_ops.placeholder(dtypes.float32, shape=input_sizes)
+      t2 = array_ops.placeholder(dtypes.float32, shape=out_backprop_sizes)
+      with self.test_scope():
+        tensor = gen_nn_ops.conv2d_backprop_filter(
+            input=t1,
+            filter_sizes=filter_sizes,
+            out_backprop=t2,
+            strides=strides,
+            dilations=dilations,
+            padding=padding,
+            data_format=data_format_dst)
+
+      value = sess.run(tensor, {t1: x1, t2: x2})
+      self.assertAllEqual(filter_sizes, value.shape)
+      self.assertAllClose(expected, value, 1e-3)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x1Filter(self, data_format):
+    expected_output = [8056, 8432, 8312, 8704, 8568, 8976]
+    self._VerifyValues(
+        input_sizes=[1, 4, 4, 3],
+        filter_sizes=[1, 1, 3, 2],
+        out_backprop_sizes=[1, 4, 4, 2],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2Filter(self, data_format):
+    expected_output = [120, 141]
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 3, 2, 1],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterDepth1(self, data_format):
+    expected_output = [5, 8, 14, 17]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Filter(self, data_format):
+    expected_output = [
+        17, 22, 27, 22, 29, 36, 27, 36, 45, 32, 43, 54, 37, 50, 63, 42, 57, 72,
+        62, 85, 108, 67, 92, 117, 72, 99, 126, 77, 106, 135, 82, 113, 144, 87,
+        120, 153
+    ]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 3],
+        filter_sizes=[2, 2, 3, 3],
+        out_backprop_sizes=[1, 1, 2, 3],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterStride3Width5(self, data_format):
+    expected_output = [9, 12]
+    self._VerifyValues(
+        input_sizes=[1, 1, 5, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[3, 3],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterStride3Width6(self, data_format):
+    expected_output = [9, 12]
+    self._VerifyValues(
+        input_sizes=[1, 1, 6, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[3, 3],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x2FilterStride3Width7(self, data_format):
+    expected_output = [9, 12]
+    self._VerifyValues(
+        input_sizes=[1, 1, 7, 1],
+        filter_sizes=[1, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[3, 3],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x3Filter(self, data_format):
+    expected_output = [5, 8, 11]
+    self._VerifyValues(
+        input_sizes=[1, 1, 4, 1],
+        filter_sizes=[1, 3, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x3FilterSame(self, data_format):
+    expected_output = [20, 30, 20]
+    self._VerifyValues(
+        input_sizes=[1, 1, 4, 1],
+        filter_sizes=[1, 3, 1, 1],
+        out_backprop_sizes=[1, 1, 4, 1],
+        strides=[1, 1],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D1x3FilterSameOutbackprop2(self, data_format):
+    expected_output = [7, 10, 3]
+    self._VerifyValues(
+        input_sizes=[1, 1, 4, 1],
+        filter_sizes=[1, 3, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[2, 2],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterC1Same(self, data_format):
+    expected_output = [91, 58, 32, 17]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 2, 3, 1],
+        strides=[1, 1],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterStride2(self, data_format):
+    expected_output = [92, 102, 112]
+    self._VerifyValues(
+        input_sizes=[1, 3, 5, 1],
+        filter_sizes=[1, 3, 1, 1],
+        out_backprop_sizes=[1, 2, 2, 1],
+        strides=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2FilterStride2Same(self, data_format):
+    expected_output = [7, 2, 16, 5]
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 2, 1],
+        strides=[2, 2],
+        padding="SAME",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Depth3ValidBackpropFilterStride1x1Dilation2x1(
+      self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 3, 6, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 5, 1],
+        strides=[1, 1],
+        dilations=[2, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[55, 70, 235, 250])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Depth1ValidBackpropFilterDilation1x2(self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 1],
+        out_backprop_sizes=[1, 1, 1, 1],
+        strides=[1, 1],
+        dilations=[1, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[1, 3, 4, 6])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DEmptyBackpropFilterDilation1x2(self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1],
+        filter_sizes=[2, 2, 1, 0],
+        out_backprop_sizes=[1, 1, 1, 0],
+        strides=[1, 1],
+        dilations=[1, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=np.zeros([0]))
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2D2x2Depth3ValidBackpropFilterDilation2x2(self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 3, 4, 3],
+        filter_sizes=[2, 2, 3, 3],
+        out_backprop_sizes=[1, 1, 2, 3],
+        strides=[1, 1],
+        dilations=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[
+            17, 22, 27, 22, 29, 36, 27, 36, 45, 47, 64, 81, 52, 71, 90, 57, 78,
+            99, 137, 190, 243, 142, 197, 252, 147, 204, 261, 167, 232, 297, 172,
+            239, 306, 177, 246, 315
+        ])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DKernelSizeMatchesInputSizeBackpropFilterDilation2x2(
+      self, data_format):
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 1],
+        filter_sizes=[2, 2, 1, 2],
+        out_backprop_sizes=[1, 1, 1, 2],
+        strides=[1, 1],
+        dilations=[2, 2],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=[1, 2, 3, 6, 7, 14, 9, 18])
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testConv2DGroupedFilter(self, data_format):
+    expected_output = [1, 4, 3, 8, 5, 12, 7, 16]
+    self._VerifyValues(
+        input_sizes=[1, 2, 2, 2],
+        filter_sizes=[2, 2, 1, 2],
+        out_backprop_sizes=[1, 1, 1, 2],
+        strides=[1, 1],
+        padding="VALID",
+        data_format_src="NHWC",
+        data_format_dst=data_format,
+        expected=expected_output)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/conv3d_test.py

@@ -0,0 +1,860 @@
+# 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.
+# ==============================================================================
+"""Tests for 3D convolutions using the XLA JIT."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import test_utils
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops import nn_ops
+import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import
+from tensorflow.python.platform import googletest
+
+
+CONV_CONFIGS = (
+    ("_Conv3D_data_format_NDHWC", "NDHWC", "Conv3D"),
+    ("_Conv3D_data_format_NCDHW", "NCDHW", "Conv3D"),
+    ("_Conv_data_format_NDHWC", "NDHWC", "Conv"),
+    ("_Conv_data_format_NCDHW", "NCDHW", "Conv"),
+)
+
+
+# Test outputs computed in prod (colab) by running nn.conv3d on a GPU device
+# with its GPU (non-xla) kernel.
+class Conv3DTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _VerifyValues(
+      self,
+      input_sizes=None,
+      filter_sizes=None,
+      strides=None,
+      dilations=None,
+      padding=None,
+      data_format_src="NDHWC",
+      data_format_dst="NDHWC",
+      expected=None,
+      op_name="Conv3D",
+  ):
+    """Tests that tf.nn.conv3d produces the expected value.
+
+    Args:
+      input_sizes: Input tensor dimensions in [batch, input_rows, input_cols,
+        input_depth].
+      filter_sizes: Filter tensor dimensions in [kernel_rows, kernel_cols,
+        input_depth, output_depth].
+      strides: Strides.
+      dilations: RHS dilations.
+      padding: Padding type.
+      data_format_src: Data format input is in.
+      data_format_dst: Data format verification will run and input is converted
+        to.
+      expected: Expected output.
+      op_name: Name of operation to test (Conv/Conv2D)
+    """
+
+    total_size_1 = np.prod(input_sizes)
+    total_size_2 = np.prod(filter_sizes)
+    x1 = np.reshape(
+        [f * 1.0 / total_size_1 for f in range(1, total_size_1 + 1)],
+        input_sizes,
+    )
+    x2 = np.reshape(
+        [f * 1.0 / total_size_2 for f in range(1, total_size_2 + 1)],
+        filter_sizes,
+    )
+    strides = [1] + strides + [1]
+    if dilations is None:
+      dilations = [1, 1, 1]
+    dilations = [1] + dilations + [1]
+
+    # Convert between data formats.
+    expected = test_utils.ConvertBetweenDataFormats(
+        expected, data_format_src, data_format_dst
+    )
+    x1 = test_utils.ConvertBetweenDataFormats(
+        x1, data_format_src, data_format_dst
+    )
+    input_sizes = test_utils.PermuteDimsBetweenDataFormats(
+        input_sizes, data_format_src, data_format_dst
+    )
+    strides = test_utils.PermuteDimsBetweenDataFormats(
+        strides, data_format_src, data_format_dst
+    )
+    dilations = test_utils.PermuteDimsBetweenDataFormats(
+        dilations, data_format_src, data_format_dst
+    )
+
+    with self.session() as sess:
+      t1 = array_ops.placeholder(dtypes.bfloat16, shape=input_sizes)
+      t2 = array_ops.placeholder(dtypes.bfloat16, shape=filter_sizes)
+      with self.test_scope():
+        if op_name == "Conv":
+          conv_format = (
+              "CHANNELS_LAST"
+              if data_format_dst == "NDHWC"
+              else "CHANNELS_FIRST"
+          )
+          out = gen_nn_ops.conv(
+              t1,
+              t2,
+              strides=strides,
+              padding=padding,
+              data_format=conv_format,
+              dilations=dilations,
+          )
+        elif op_name == "Conv3D":
+          out = nn_ops.conv3d(
+              t1,
+              t2,
+              strides=strides,
+              padding=padding,
+              data_format=data_format_dst,
+              dilations=dilations,
+          )
+        else:
+          raise ValueError("Invalid op name: %s" % op_name)
+
+      value = sess.run(out, {t1: x1, t2: x2})
+      self.assertAllCloseAccordingToType(expected, value)
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3D1x1x1Filter(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            0.18518518518518517,
+            0.2222222222222222,
+            0.25925925925925924,
+            0.4074074074074074,
+            0.5,
+            0.5925925925925926,
+            0.6296296296296297,
+            0.7777777777777777,
+            0.9259259259259259,
+            0.8518518518518519,
+            1.0555555555555556,
+            1.259259259259259,
+            1.074074074074074,
+            1.3333333333333333,
+            1.5925925925925926,
+            1.2962962962962963,
+            1.6111111111111112,
+            1.9259259259259258,
+        ],
+        [1, 2, 3, 1, 3],
+    )
+
+    # These are equivalent to the Conv2D1x1 case.
+    self._VerifyValues(
+        input_sizes=[1, 2, 3, 1, 3],
+        filter_sizes=[1, 1, 1, 3, 3],
+        strides=[1, 1, 1],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+    self._VerifyValues(
+        input_sizes=[1, 2, 1, 3, 3],
+        filter_sizes=[1, 1, 1, 3, 3],
+        strides=[1, 1, 1],
+        padding="VALID",
+        expected=np.reshape(expected_output, [1, 2, 1, 3, 3]),
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+    self._VerifyValues(
+        input_sizes=[1, 1, 2, 3, 3],
+        filter_sizes=[1, 1, 1, 3, 3],
+        strides=[1, 1, 1],
+        padding="VALID",
+        expected=np.reshape(expected_output, [1, 1, 2, 3, 3]),
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3D1x1x1Filter2x1x1Dilation(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            0.05555555555555555,
+            0.1111111111111111,
+            0.16666666666666666,
+            0.2222222222222222,
+            0.2777777777777778,
+            0.3333333333333333,
+            0.3888888888888889,
+            0.4444444444444444,
+            0.5,
+            0.5555555555555556,
+            0.6111111111111112,
+            0.6666666666666666,
+            0.7222222222222222,
+            0.7777777777777778,
+            0.8333333333333334,
+            0.8888888888888888,
+            0.9444444444444444,
+            1.0,
+        ],
+        [1, 3, 6, 1, 1],
+    )
+
+    self._VerifyValues(
+        input_sizes=[1, 3, 6, 1, 1],
+        filter_sizes=[1, 1, 1, 1, 1],
+        strides=[1, 1, 1],
+        padding="VALID",
+        dilations=[2, 1, 1],
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  # Expected values computed using scipy's correlate function.
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3D2x2x2Filter(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            3.7719907407407405,
+            3.850694444444445,
+            3.929398148148149,
+            4.265046296296295,
+            4.357638888888888,
+            4.450231481481481,
+            6.730324074074074,
+            6.892361111111109,
+            7.054398148148148,
+            7.223379629629629,
+            7.399305555555557,
+            7.575231481481481,
+            9.688657407407408,
+            9.934027777777779,
+            10.17939814814815,
+            10.181712962962962,
+            10.440972222222221,
+            10.700231481481481,
+        ],
+        [1, 3, 1, 2, 3],
+    )
+    # expected_shape = [1, 3, 1, 2, 5]
+    self._VerifyValues(
+        input_sizes=[1, 4, 2, 3, 3],  # b, z, y, x, fin
+        filter_sizes=[2, 2, 2, 3, 3],  # z, y, x, fin, fout
+        strides=[1, 1, 1],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3D2x2x2Filter1x2x1Dilation(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            1.1388888888888888,
+            1.2013888888888888,
+            1.3263888888888888,
+            1.3888888888888888,
+            1.5138888888888888,
+            1.5763888888888888,
+            1.701388888888889,
+            1.763888888888889,
+            2.263888888888889,
+            2.3263888888888893,
+            2.451388888888889,
+            2.513888888888889,
+            2.6388888888888893,
+            2.701388888888889,
+            2.826388888888889,
+            2.888888888888889,
+            3.388888888888889,
+            3.451388888888889,
+            3.576388888888889,
+            3.6388888888888884,
+            3.7638888888888893,
+            3.8263888888888893,
+            3.9513888888888893,
+            4.013888888888889,
+        ],
+        [1, 3, 4, 2, 1],
+    )
+
+    self._VerifyValues(
+        input_sizes=[1, 4, 6, 3, 1],
+        filter_sizes=[2, 2, 2, 1, 1],
+        strides=[1, 1, 1],
+        padding="VALID",
+        dilations=[1, 2, 1],
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3DStrides(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            0.06071428571428571,
+            0.08988095238095238,
+            0.10238095238095238,
+            0.11488095238095238,
+            0.12738095238095237,
+            0.13988095238095238,
+            0.08452380952380953,
+            0.26071428571428573,
+            0.35238095238095235,
+            0.36488095238095236,
+            0.3773809523809524,
+            0.3898809523809524,
+            0.4023809523809524,
+            0.23452380952380952,
+            0.46071428571428574,
+            0.6148809523809524,
+            0.6273809523809524,
+            0.6398809523809523,
+            0.6523809523809524,
+            0.6648809523809525,
+            0.3845238095238095,
+            1.1273809523809524,
+            1.4898809523809524,
+            1.5023809523809524,
+            1.5148809523809523,
+            1.5273809523809523,
+            1.5398809523809525,
+            0.8845238095238095,
+            1.3273809523809526,
+            1.7523809523809522,
+            1.764880952380952,
+            1.7773809523809523,
+            1.7898809523809525,
+            1.8023809523809526,
+            1.0345238095238096,
+            1.5273809523809525,
+            2.0148809523809526,
+            2.0273809523809523,
+            2.0398809523809525,
+            2.052380952380952,
+            2.0648809523809524,
+            1.1845238095238095,
+            2.1940476190476192,
+            2.8898809523809526,
+            2.9023809523809527,
+            2.9148809523809525,
+            2.9273809523809526,
+            2.9398809523809524,
+            1.6845238095238095,
+            2.394047619047619,
+            3.1523809523809523,
+            3.1648809523809525,
+            3.177380952380952,
+            3.1898809523809524,
+            3.2023809523809526,
+            1.8345238095238097,
+            2.594047619047619,
+            3.4148809523809525,
+            3.427380952380952,
+            3.4398809523809524,
+            3.4523809523809526,
+            3.4648809523809523,
+            1.9845238095238096,
+        ],
+        [1, 3, 3, 7, 1],
+    )
+    self._VerifyValues(
+        input_sizes=[1, 5, 8, 7, 1],
+        filter_sizes=[1, 2, 3, 1, 1],
+        strides=[2, 3, 1],  # different stride for each spatial dimension
+        padding="SAME",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3D2x2x2FilterStride2(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            3.7719907407407405,
+            3.850694444444445,
+            3.929398148148149,
+            9.688657407407408,
+            9.934027777777779,
+            10.17939814814815,
+        ],
+        [1, 2, 1, 1, 3],
+    )
+    self._VerifyValues(
+        input_sizes=[1, 4, 2, 3, 3],
+        filter_sizes=[2, 2, 2, 3, 3],
+        strides=[2, 2, 2],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3DStride3(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            1.5114087301587302,
+            1.5716765873015872,
+            1.6319444444444446,
+            1.5634920634920635,
+            1.6267361111111112,
+            1.6899801587301588,
+            1.6155753968253967,
+            1.681795634920635,
+            1.748015873015873,
+            1.9280753968253967,
+            2.012152777777778,
+            2.096230158730159,
+            1.9801587301587302,
+            2.067212301587302,
+            2.154265873015873,
+            2.0322420634920637,
+            2.122271825396825,
+            2.2123015873015874,
+            4.428075396825396,
+            4.65500992063492,
+            4.881944444444444,
+            4.480158730158729,
+            4.710069444444444,
+            4.939980158730158,
+            4.532242063492063,
+            4.7651289682539675,
+            4.9980158730158735,
+            4.844742063492064,
+            5.095486111111112,
+            5.346230158730158,
+            4.896825396825397,
+            5.150545634920635,
+            5.4042658730158735,
+            4.94890873015873,
+            5.205605158730158,
+            5.462301587301588,
+        ],
+        [1, 2, 2, 3, 3],
+    )
+    self._VerifyValues(
+        input_sizes=[1, 6, 7, 8, 2],
+        filter_sizes=[3, 2, 1, 2, 3],
+        strides=[3, 3, 3],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testConv3D2x2x2FilterStride2Same(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            3.7719907407407405,
+            3.850694444444445,
+            3.929398148148149,
+            2.0162037037037037,
+            2.0659722222222223,
+            2.1157407407407405,
+            9.688657407407408,
+            9.934027777777779,
+            10.17939814814815,
+            4.599537037037037,
+            4.732638888888889,
+            4.8657407407407405,
+        ],
+        [1, 2, 1, 2, 3],
+    )
+    self._VerifyValues(
+        input_sizes=[1, 4, 2, 3, 3],
+        filter_sizes=[2, 2, 2, 3, 3],
+        strides=[2, 2, 2],
+        padding="SAME",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testKernelSmallerThanStride(self, data_format, op_name):
+    expected_output = np.reshape(
+        [
+            0.037037037037037035,
+            0.1111111111111111,
+            0.25925925925925924,
+            0.3333333333333333,
+            0.7037037037037037,
+            0.7777777777777778,
+            0.9259259259259259,
+            1.0,
+        ],
+        [1, 2, 2, 2, 1],
+    )
+
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 3, 1],
+        filter_sizes=[1, 1, 1, 1, 1],
+        strides=[2, 2, 2],
+        padding="SAME",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+    self._VerifyValues(
+        input_sizes=[1, 3, 3, 3, 1],
+        filter_sizes=[1, 1, 1, 1, 1],
+        strides=[2, 2, 2],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+    expected_output = np.reshape(
+        [
+            0.5408163265306123,
+            0.5801749271137027,
+            0.28061224489795916,
+            0.8163265306122448,
+            0.8556851311953353,
+            0.4030612244897959,
+            0.41873177842565595,
+            0.43403790087463556,
+            0.19642857142857142,
+            2.4693877551020407,
+            2.5087463556851315,
+            1.1377551020408163,
+            2.7448979591836733,
+            2.7842565597667637,
+            1.260204081632653,
+            1.168731778425656,
+            1.1840379008746356,
+            0.5178571428571429,
+            1.0951166180758019,
+            1.1060495626822158,
+            0.4464285714285714,
+            1.1716472303206997,
+            1.1825801749271136,
+            0.4770408163265306,
+            0.3691690962099125,
+            0.37244897959183676,
+            0.125,
+        ],
+        [1, 3, 3, 3, 1],
+    )
+    self._VerifyValues(
+        input_sizes=[1, 7, 7, 7, 1],
+        filter_sizes=[2, 2, 2, 1, 1],
+        strides=[3, 3, 3],
+        padding="SAME",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+    expected_output = np.reshape(
+        [
+            0.5408163265306123,
+            0.5801749271137027,
+            0.8163265306122448,
+            0.8556851311953353,
+            2.4693877551020407,
+            2.5087463556851315,
+            2.7448979591836733,
+            2.7842565597667637,
+        ],
+        [1, 2, 2, 2, 1],
+    )
+    self._VerifyValues(
+        input_sizes=[1, 7, 7, 7, 1],
+        filter_sizes=[2, 2, 2, 1, 1],
+        strides=[3, 3, 3],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  @parameterized.named_parameters(*CONV_CONFIGS)
+  def testKernelSizeMatchesInputSize(self, data_format, op_name):
+    expected_output = np.reshape([1.5625, 1.875], [1, 1, 1, 1, 2])
+    self._VerifyValues(
+        input_sizes=[1, 2, 1, 2, 1],
+        filter_sizes=[2, 1, 2, 1, 2],
+        strides=[1, 1, 1],
+        padding="VALID",
+        expected=expected_output,
+        data_format_src="NDHWC",
+        data_format_dst=data_format,
+        op_name=op_name,
+    )
+
+  def testConvExpandedBatch(self):
+    tensor_in_sizes_batch = [10, 2, 3, 1, 3]
+    tensor_in_sizes_expanded_batch = [2, 5, 2, 3, 1, 3]
+    batch_dims = 2
+    filter_in_sizes = [1, 1, 1, 3, 3]
+    filter_in = np.arange(
+        1, np.prod(filter_in_sizes) + 1, dtype=np.float32
+    ).reshape(filter_in_sizes)
+    x1 = np.arange(
+        1, np.prod(tensor_in_sizes_batch) + 1, dtype=np.float32
+    ).reshape(tensor_in_sizes_batch)
+    x2 = x1.reshape(tensor_in_sizes_expanded_batch)
+
+    with self.session() as sess:
+      t1 = array_ops.placeholder(dtypes.bfloat16, shape=tensor_in_sizes_batch)
+      t2 = array_ops.placeholder(
+          dtypes.bfloat16, shape=tensor_in_sizes_expanded_batch
+      )
+      filter_t = array_ops.placeholder(dtypes.bfloat16, shape=filter_in_sizes)
+
+      out1 = gen_nn_ops.conv(
+          t1, filter_t, strides=[1, 1, 1, 1, 1], padding="VALID"
+      )
+      out2 = gen_nn_ops.conv(
+          t2,
+          filter_t,
+          strides=[1, 1, 1, 1, 1],
+          padding="VALID",
+          batch_dims=batch_dims,
+      )
+      value1 = sess.run(out1, {t1: x1, filter_t: filter_in})
+      value2 = sess.run(out2, {t2: x2, filter_t: filter_in})
+
+      self.assertEqual(list(value1.shape), tensor_in_sizes_batch)
+      self.assertEqual(list(value2.shape), tensor_in_sizes_expanded_batch)
+      self.assertAllCloseAccordingToType(value1, value2.reshape(value1.shape))
+
+
+# Test cloned from
+# tensorflow/python/kernel_tests/conv3d_backprop_filter_v2_grad_test.py
+class Conv3DBackpropFilterV2GradTest(xla_test.XLATestCase):
+
+  def testGradient(self):
+    with self.session(), self.test_scope():
+      for padding in ["SAME", "VALID"]:
+        for stride in [1, 2]:
+          np.random.seed(1)
+          in_shape = [2, 4, 3, 3, 2]
+          in_val = constant_op.constant(
+              2 * np.random.random_sample(in_shape) - 1, dtype=dtypes.float32)
+          filter_shape = [3, 3, 3, 2, 3]
+          strides = [1, stride, stride, stride, 1]
+          # Make a convolution op with the current settings, just to easily get
+          # the shape of the output.
+          conv_out = nn_ops.conv3d(in_val,
+                                   array_ops.zeros(filter_shape), strides,
+                                   padding)
+          out_backprop_shape = conv_out.get_shape().as_list()
+          out_backprop_val = constant_op.constant(
+              2 * np.random.random_sample(out_backprop_shape) - 1,
+              dtype=dtypes.float32)
+          output = nn_ops.conv3d_backprop_filter_v2(in_val, filter_shape,
+                                                    out_backprop_val, strides,
+                                                    padding)
+          err = gradient_checker.compute_gradient_error(
+              [in_val, out_backprop_val], [in_shape, out_backprop_shape],
+              output, filter_shape)
+          print("conv3d_backprop_filter gradient err = %g " % err)
+          err_tolerance = 1e-3
+          self.assertLess(err, err_tolerance)
+
+
+# Test cloned from tensorflow/python/kernel_tests/conv3d_transpose_test.py
+class Conv3DTransposeTest(xla_test.XLATestCase):
+
+  def testConv3DTransposeSingleStride(self):
+    with self.session(), self.test_scope():
+      strides = [1, 1, 1, 1, 1]
+
+      # Input, output: [batch, depth, height, width, channel]
+      x_shape = [2, 5, 6, 4, 3]
+      y_shape = [2, 5, 6, 4, 2]
+
+      # Filter: [kernel_depth, kernel_height, kernel_width, out_depth, in_depth]
+      f_shape = [3, 3, 3, 2, 3]
+
+      x = constant_op.constant(
+          1.0, shape=x_shape, name="x", dtype=dtypes.float32)
+      f = constant_op.constant(
+          1.0, shape=f_shape, name="filter", dtype=dtypes.float32)
+      output = nn_ops.conv3d_transpose(
+          x, f, y_shape, strides=strides, padding="SAME")
+      value = self.evaluate(output)
+
+      # We count the number of cells being added at the locations in the output.
+      # At the center, #cells = kernel_depth * kernel_height * kernel_width
+      # At the corners, #cells = ceil(kernel_depth/2) * ceil(kernel_height/2)
+      #                          * ceil(kernel_width/2)
+      # At the edges, #cells =
+      #   kernel_depth * ceil(kernel_height/2) * ceil(kernel_width/2) or
+      #   ceil(kernel_depth/2) * kernel_height * ceil(kernel_width/2) or
+      #   ceil(kernel_depth/2) * ceil(kernel_height/2) * kernel_width
+      # At the borders, #cells =
+      #   ceil(kernel_depth/2) * kernel_height * kernel_width or
+      #   kernel_depth * ceil(kernel_height/2) * kernel_width or
+      #   kernel_depth * kernel_height * ceil(kernel_width/2)
+
+      for n in range(x_shape[0]):
+        for k in range(f_shape[3]):
+          for w in range(y_shape[3]):
+            for h in range(y_shape[2]):
+              for d in range(y_shape[1]):
+                d_in = d > 0 and d < y_shape[1] - 1
+                h_in = h > 0 and h < y_shape[2] - 1
+                w_in = w > 0 and w < y_shape[3] - 1
+                if d_in + h_in + w_in == 3:
+                  target = 27 * 3.0
+                elif d_in + h_in + w_in == 2:
+                  target = 18 * 3.0
+                elif d_in or h_in or w_in:
+                  target = 12 * 3.0
+                else:
+                  target = 8 * 3.0
+                self.assertAllClose(target, value[n, d, h, w, k])
+
+  def testConv3DTransposeSame(self):
+    with self.session(), self.test_scope():
+      strides = [1, 2, 2, 2, 1]
+
+      # Input, output: [batch, depth, height, width, depth]
+      x_shape = [2, 5, 6, 4, 3]
+      y_shape = [2, 10, 12, 8, 2]
+
+      # Filter: [kernel_depth, kernel_height, kernel_width, out_depth, in_depth]
+      f_shape = [3, 3, 3, 2, 3]
+
+      x = constant_op.constant(
+          1.0, shape=x_shape, name="x", dtype=dtypes.float32)
+      f = constant_op.constant(
+          1.0, shape=f_shape, name="filter", dtype=dtypes.float32)
+      output = nn_ops.conv3d_transpose(
+          x, f, y_shape, strides=strides, padding="SAME")
+      value = self.evaluate(output)
+
+      for n in range(x_shape[0]):
+        for k in range(f_shape[3]):
+          for w in range(y_shape[3]):
+            for h in range(y_shape[2]):
+              for d in range(y_shape[1]):
+                # We add a case for locations divisible by the stride.
+                d_in = d % strides[1] == 0 and 0 < d < y_shape[1] - 1
+                h_in = h % strides[2] == 0 and 0 < h < y_shape[2] - 1
+                w_in = w % strides[3] == 0 and 0 < w < y_shape[3] - 1
+                if d_in + h_in + w_in == 3:
+                  target = 8 * 3.0
+                elif d_in + h_in + w_in == 2:
+                  target = 4 * 3.0
+                elif d_in or h_in or w_in:
+                  target = 2 * 3.0
+                else:
+                  target = 3.0
+                self.assertAllClose(target, value[n, d, h, w, k])
+
+  def testConv3DTransposeValid(self):
+    with self.session(), self.test_scope():
+      strides = [1, 2, 2, 2, 1]
+
+      # Input, output: [batch, depth, height, width, depth]
+      x_shape = [2, 5, 6, 4, 3]
+      y_shape = [2, 11, 13, 9, 2]
+
+      # Filter: [kernel_depth, kernel_height, kernel_width, out_depth, in_depth]
+      f_shape = [3, 3, 3, 2, 3]
+
+      x = constant_op.constant(
+          1.0, shape=x_shape, name="x", dtype=dtypes.float32)
+      f = constant_op.constant(
+          1.0, shape=f_shape, name="filter", dtype=dtypes.float32)
+      output = nn_ops.conv3d_transpose(
+          x, f, y_shape, strides=strides, padding="VALID")
+      value = self.evaluate(output)
+
+      cache_values = np.zeros(y_shape, dtype=np.float32)
+
+      # The amount of padding added
+      pad = 1
+
+      for n in range(x_shape[0]):
+        for k in range(f_shape[3]):
+          for w in range(y_shape[3]):
+            for h in range(y_shape[2]):
+              for d in range(y_shape[1]):
+                # We add a case for locations divisible by the stride.
+                d_in = d % strides[1] == 0 and pad < d < y_shape[1] - 1 - pad
+                h_in = h % strides[2] == 0 and pad < h < y_shape[2] - 1 - pad
+                w_in = w % strides[3] == 0 and pad < w < y_shape[3] - 1 - pad
+                if d_in + h_in + w_in == 3:
+                  target = 8 * 3.0
+                elif d_in + h_in + w_in == 2:
+                  target = 4 * 3.0
+                elif d_in or h_in or w_in:
+                  target = 2 * 3.0
+                else:
+                  target = 3.0
+                cache_values[n, d, h, w, k] = target
+
+          # copy values in the border
+          cache_values[n, :, :, 0, k] = cache_values[n, :, :, 1, k]
+          cache_values[n, :, :, -1, k] = cache_values[n, :, :, -2, k]
+          cache_values[n, :, 0, :, k] = cache_values[n, :, 1, :, k]
+          cache_values[n, :, -1, :, k] = cache_values[n, :, -2, :, k]
+          cache_values[n, 0, :, :, k] = cache_values[n, 1, :, :, k]
+          cache_values[n, -1, :, :, k] = cache_values[n, -2, :, :, k]
+
+    self.assertAllClose(cache_values, value)
+
+  def testGradient(self):
+    x_shape = [2, 3, 4, 3, 2]
+    f_shape = [3, 3, 3, 2, 2]
+    y_shape = [2, 6, 8, 6, 2]
+    strides = [1, 2, 2, 2, 1]
+    np.random.seed(1)  # Make it reproducible.
+    x_val = np.random.random_sample(x_shape).astype(np.float64)
+    f_val = np.random.random_sample(f_shape).astype(np.float64)
+    with self.session(), self.test_scope():
+      x = constant_op.constant(x_val, name="x", dtype=dtypes.float32)
+      f = constant_op.constant(f_val, name="f", dtype=dtypes.float32)
+      output = nn_ops.conv3d_transpose(
+          x, f, y_shape, strides=strides, padding="SAME")
+      err = gradient_checker.compute_gradient_error([x, f], [x_shape, f_shape],
+                                                    output, y_shape)
+    print("conv3d_transpose gradient err = %g " % err)
+    err_tolerance = 0.001
+    self.assertLess(err, err_tolerance)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/conv_node_name_test.py

@@ -0,0 +1,118 @@
+# 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.
+# ==============================================================================
+"""Tests for Convolution node name match via the XLA JIT.
+
+The canned results in these tests are created by running each test using the
+Tensorflow CPU device and saving the output.
+"""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import ops
+from tensorflow.python.layers import layers
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.platform import googletest
+
+
+class ConvolutionNodeNameTest(xla_test.XLATestCase):
+  """Verify convolution node name match.
+
+  Verify convolution node names on TPU and CPU match with dilation > 1.
+  """
+
+  def _verifyNodeNameMatch(self, layer, input_sizes, filter_sizes, strides,
+                           dilations):
+
+    def _GetNodeNames(use_xla):
+      with self.session():
+        input_tensor = array_ops.placeholder(np.float32, shape=input_sizes)
+
+        if use_xla:
+          with self.device_scope():
+            # pylint: disable=protected-access
+            graph = ops.get_default_graph()
+            graph._set_control_flow_context(
+                control_flow_ops.XLAControlFlowContext())
+            # pylint: enable=protected-access
+            conv2d_op = layer(
+                filters=64,
+                kernel_size=filter_sizes,
+                dilation_rate=dilations,
+                padding="same")
+            _ = conv2d_op(input_tensor)
+            return [n.name for n in ops.get_default_graph().as_graph_def().node]
+        else:
+          with ops.device("CPU"):
+            conv2d_op = layer(
+                filters=64,
+                kernel_size=filter_sizes,
+                dilation_rate=dilations,
+                padding="same")
+            _ = conv2d_op(input_tensor)
+            names = [
+                n.name for n in ops.get_default_graph().as_graph_def().node
+            ]
+            # filter out space to depth ops.
+            return [
+                name for name in names
+                if "space" not in name and "Space" not in name
+            ]
+
+    xla_names = _GetNodeNames(use_xla=True)
+    no_xla_names = _GetNodeNames(use_xla=False)
+
+    # CPU path creates some additional nodes to handle dilations.
+    # TODO(b/138804006): Remove this when CPU & GPU support dilations.
+    filtered_no_xla_names = []
+    for name in no_xla_names:
+      if ("dilation_rate" in name or "filter_shape" in name or "stack" in name):
+        continue
+      else:
+        filtered_no_xla_names.append(name)
+
+    self.assertListEqual(xla_names, filtered_no_xla_names)
+
+  def testConv1DNodeNameMatch(self):
+    input_sizes = [8, 16, 3]
+    filter_sizes = [7]
+    strides = 1
+    dilations = [2]
+    layer = layers.Conv1D
+    self._verifyNodeNameMatch(layer, input_sizes, filter_sizes, strides,
+                              dilations)
+
+  def testConv2DNodeNameMatch(self):
+    input_sizes = [8, 16, 16, 3]
+    filter_sizes = [7, 7]
+    strides = 1
+    dilations = [2, 2]
+    layer = layers.Conv2D
+    self._verifyNodeNameMatch(layer, input_sizes, filter_sizes, strides,
+                              dilations)
+
+  def testConv3DNodeNameMatch(self):
+    input_sizes = [8, 16, 16, 16, 3]
+    filter_sizes = [7, 7, 7]
+    strides = 1
+    dilations = [2, 2, 2]
+    layer = layers.Conv3D
+    self._verifyNodeNameMatch(layer, input_sizes, filter_sizes, strides,
+                              dilations)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/data_format_ops_test.py

@@ -0,0 +1,147 @@
+# Copyright 2018 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 the DataFormatVecPermute operator."""
+
+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 nn_ops
+from tensorflow.python.platform import test
+
+
+class XlaDataFormatDimMapTest(xla_test.XLATestCase):
+
+  def _test(self, input_data, src_format, dst_format, expected):
+    for dtype in {np.int32, np.int64}:
+      x = np.array(input_data, dtype=dtype)
+      with self.session() as session:
+        with self.test_scope():
+          placeholder = array_ops.placeholder(dtypes.as_dtype(x.dtype), x.shape)
+          param = {placeholder: x}
+          output = nn_ops.data_format_dim_map(
+              placeholder, src_format=src_format, dst_format=dst_format)
+        result = session.run(output, param)
+      self.assertAllEqual(result, expected)
+
+  def test(self):
+    self._test(0, "NHWC", "NCHW", 0)
+    self._test(1, "NHWC", "NCHW", 2)
+    self._test(2, "NHWC", "NCHW", 3)
+    self._test(3, "NHWC", "NCHW", 1)
+    self._test(-1, "NHWC", "NCHW", 1)
+    self._test(-2, "NHWC", "NCHW", 3)
+    self._test(-3, "NHWC", "NCHW", 2)
+    self._test(-4, "NHWC", "NCHW", 0)
+    self._test([1, 3], "NHWC", "NCHW", [2, 1])
+    self._test([1, 3, -2], "NHWC", "NCHW", [2, 1, 3])
+    self._test([1, -3, -2], "NHWC", "NCHW", [2, 2, 3])
+    self._test([[1, -3], [1, -1]], "NHWC", "NCHW", [[2, 2], [2, 1]])
+
+    self._test([1, -3, -2], "NHWC", "NCHW", [2, 2, 3])
+    self._test([-4, -3, -2, -1, 0, 1, 2, 3], "NHWC", "HWNC",
+               [2, 0, 1, 3, 2, 0, 1, 3])
+    self._test([-4, -3, -2, -1, 0, 1, 2, 3], "NHWC", "WHCN",
+               [3, 1, 0, 2, 3, 1, 0, 2])
+    self._test([-4, -3, -2, -1, 0, 1, 2, 3], "qwer", "rewq",
+               [3, 2, 1, 0, 3, 2, 1, 0])
+
+    self._test(0, "NDHWC", "NCDHW", 0)
+    self._test(1, "NDHWC", "NCDHW", 2)
+    self._test(2, "NDHWC", "NCDHW", 3)
+    self._test(3, "NDHWC", "NCDHW", 4)
+    self._test(4, "NDHWC", "NCDHW", 1)
+    self._test([1, 4], "NDHWC", "NCDHW", [2, 1])
+    self._test([1, 4, -2], "NDHWC", "NCDHW", [2, 1, 4])
+    self._test([1, -3, -2], "NDHWC", "NCDHW", [2, 3, 4])
+    self._test([[1, -4], [1, -1]], "NDHWC", "NCDHW", [[2, 2], [2, 1]])
+
+    self._test([1, -3, -2], "NDHWC", "NCDHW", [2, 3, 4])
+    self._test([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4], "NDHWC", "DHWNC",
+               [3, 0, 1, 2, 4, 3, 0, 1, 2, 4])
+    self._test([-5, -4, -3, -2, -1, 0, 1, 2, 3, 4], "NDHWC", "WHDCN",
+               [4, 2, 1, 0, 3, 4, 2, 1, 0, 3])
+
+
+class XlaPermuteOpTest(xla_test.XLATestCase):
+
+  def _runPermuteAndCompare(self, x, src_format, dst_format, expected):
+    with self.session() as session:
+      with self.test_scope():
+        placeholder = array_ops.placeholder(dtypes.as_dtype(x.dtype), x.shape)
+        param = {placeholder: x}
+        output = nn_ops.data_format_vec_permute(
+            placeholder, src_format=src_format, dst_format=dst_format)
+      result = session.run(output, param)
+    self.assertAllEqual(result, expected)
+
+  def testNHWCToNCHW(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([7, 4, 9, 3], dtype=dtype)
+      self._runPermuteAndCompare(x, "NHWC", "NCHW", [7, 3, 4, 9])
+
+  def testNHWCToNCHW_Size2(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([4, 9], dtype=dtype)
+      self._runPermuteAndCompare(x, "NHWC", "NCHW", [4, 9])
+
+  def testNCHWToNHWC(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([7, 4, 9, 3], dtype=dtype)
+      self._runPermuteAndCompare(x, "NCHW", "NHWC", [7, 9, 3, 4])
+
+  def testNCHWToNHWC_Size2(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([9, 3], dtype=dtype)
+      self._runPermuteAndCompare(x, "NCHW", "NHWC", [9, 3])
+
+  def testNHWCToHWNC(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([7, 4, 9, 3], dtype=dtype)
+      self._runPermuteAndCompare(x, "NHWC", "HWNC", [4, 9, 7, 3])
+
+  def testHWNCToNHWC(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([7, 4, 9, 3], dtype=dtype)
+      self._runPermuteAndCompare(x, "HWNC", "NHWC", [9, 7, 4, 3])
+
+  def testNHWCToNCHW2D(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([[7, 4], [9, 3], [4, 5], [5, 1]], dtype=dtype)
+      self._runPermuteAndCompare(x, "NHWC", "NCHW",
+                                 [[7, 4], [5, 1], [9, 3], [4, 5]])
+
+  def testNHWCToHWNC2D(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([[7, 4], [9, 3], [4, 5], [5, 1]], dtype=dtype)
+      self._runPermuteAndCompare(x, "NHWC", "HWNC",
+                                 [[9, 3], [4, 5], [7, 4], [5, 1]])
+
+  def testHWNCToNHWC2D(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([[7, 4], [9, 3], [4, 5], [5, 1]], dtype=dtype)
+      self._runPermuteAndCompare(x, "HWNC", "NHWC",
+                                 [[4, 5], [7, 4], [9, 3], [5, 1]])
+
+  def testNCHWToNHWC2D(self):
+    for dtype in {np.int32, np.int64}:
+      x = np.array([[7, 4], [9, 3], [4, 5], [5, 1]], dtype=dtype)
+      self._runPermuteAndCompare(x, "NCHW", "NHWC",
+                                 [[7, 4], [4, 5], [5, 1], [9, 3]])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/dense_layer_test.py

@@ -0,0 +1,141 @@
+# Copyright 2018 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 DenseLayer JIT compilation on the CPU and GPU devices."""
+
+import os
+
+import numpy as np
+
+from tensorflow.compiler.tests import test_utils
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.compiler.xla import jit
+from tensorflow.python.framework import ops
+from tensorflow.python.layers import layers
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+jit_scope = jit.experimental_jit_scope
+
+def GetRunMetadataLabels(run_metadata):
+  """Returns all labels in run_metadata."""
+  labels = []
+  for dev_stats in run_metadata.step_stats.dev_stats:
+    for node_stats in dev_stats.node_stats:
+      labels.append(node_stats.timeline_label)
+  return labels
+
+
+def InLabels(labels, substr):
+  """Returns true iff one of the labels contains substr."""
+  return any(substr in x for x in labels)
+
+
+class DenseLayerTest(test.TestCase):
+
+  def countXlaOps(self, labels):
+    """Count how many XlaCompile/XlaRun labels are present."""
+    xla_compile_count = sum("XlaCompile(" in x for x in labels)
+    xla_run_count = sum("XlaRun(" in x for x in labels)
+    self.assertEqual(xla_compile_count, xla_run_count)
+    return xla_run_count
+
+
+  def testDenseLayerAutoJit(self):
+    """Tests dense layer compilation in auto-jit mode.
+
+    Dense layer should be compiled into a single XlaCompile/XlaRun op pair in
+    auto-jit mode.
+    """
+
+    os.environ["TF_XLA_FLAGS"] = (
+        "--tf_xla_cpu_global_jit " + os.environ.get("TF_XLA_FLAGS", ""))
+    config = config_pb2.ConfigProto()
+    config.graph_options.optimizer_options.global_jit_level = (
+        config_pb2.OptimizerOptions.ON_1)
+
+    with self.session(config=config) as sess:
+      x = array_ops.placeholder(shape=[None, None, 3], dtype=np.float32)
+      y = layers.dense(x, 3)
+
+      self.evaluate(variables.global_variables_initializer())
+      run_metadata = config_pb2.RunMetadata()
+      test_utils.RunWithWarmup(
+          sess,
+          y, {x: np.array([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])},
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+
+    labels = GetRunMetadataLabels(run_metadata)
+    self.assertEqual(1, self.countXlaOps(labels))
+    self.assertFalse(InLabels(labels, "MatMult"))
+
+  def testDenseLayerJitScopeDefinedShape(self):
+    """Tests that the dense layer node is properly compiled in jit scope.
+
+    Dense layer with static shape input tensor should be compiled into a single
+    XlaCompile/XlaRun op pair by XLA.
+    """
+
+    with self.session() as sess:
+      x = array_ops.placeholder(shape=[2, 2, 3], dtype=np.float32)
+      with jit_scope():
+        y = layers.dense(x, 3)
+
+      self.evaluate(variables.global_variables_initializer())
+      run_metadata = config_pb2.RunMetadata()
+      test_utils.RunWithWarmup(
+          sess,
+          y, {x: np.array([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])},
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+
+    labels = GetRunMetadataLabels(run_metadata)
+    self.assertEqual(1, self.countXlaOps(labels))
+    # No need to check whether ListDiff is compiled or not because ListDiff op
+    # is not used when input tensor shape is fully defined.
+
+  def testDenseLayerJitScopeUndefinedShape(self):
+    """Tests that the dense layer node is properly compiled in jit scope.
+    """
+
+    with self.session() as sess:
+      x = array_ops.placeholder(shape=[None, None, 3], dtype=np.float32)
+      with jit_scope():
+        y = layers.dense(x, 3)
+
+      self.evaluate(variables.global_variables_initializer())
+      run_metadata = config_pb2.RunMetadata()
+      test_utils.RunWithWarmup(
+          sess,
+          y, {x: np.array([[[1, 2, 3], [4, 5, 6]], [[1, 2, 3], [4, 5, 6]]])},
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+
+    labels = GetRunMetadataLabels(run_metadata)
+    self.assertEqual(1, self.countXlaOps(labels))
+    self.assertFalse(InLabels(labels, "MatMult"))
+
+
+if __name__ == "__main__":
+  os.environ["TF_XLA_FLAGS"] = ("--tf_xla_enable_lazy_compilation=true " +
+                                os.environ.get("TF_XLA_FLAGS", ""))
+  # This test is using Tensorflow sessions which are not compatible with eager
+  # mode.
+  ops.disable_eager_execution()
+  test.main()

tensorflow/compiler/tests/depthwise_conv_op_test.py

@@ -0,0 +1,701 @@
+# Copyright 2015 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 depthwise convolutional operations."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import nn_impl
+from tensorflow.python.ops import nn_ops
+import tensorflow.python.ops.nn_grad  # pylint: disable=unused-import
+from tensorflow.python.platform import test
+
+
+# Reference implementation of depthwise_conv2d
+def ReferenceDepthwiseConv2D(input_tensor, filter_tensor, strides, padding,
+                             data_format=None):
+  # Reference implementation of depthwise convolution that uses regular
+  # convolution.
+  convs = []
+  in_channels = filter_tensor.shape[2]
+  # Use a custom implementation of depthwise conv2d using slicing.
+  for channel in range(in_channels):
+    # Slice the input along channel
+    if data_format == "NCHW":
+      input_slice = input_tensor[:, channel:channel+1, :, :]
+    else:
+      input_slice = input_tensor[:, :, :, channel:channel+1]
+
+    # Slice the filters.  Filters are  H, W, InC, DepthMultiplier
+    filter_slice = filter_tensor[:, :, channel:channel+1, :]
+    # Do conv
+    convs.append(nn_ops.conv2d(input_slice, filter_slice,
+                               strides, padding,
+                               data_format=data_format,
+                               name="depthwise_slice_%d" % channel))
+
+  # Concat along dimension.
+  if data_format == "NCHW":
+    return array_ops.concat(convs, 1)
+  else:
+    return array_ops.concat(convs, 3)
+
+
+def ConfigsToTest():
+  """Iterator for different convolution shapes, strides and paddings.
+
+  Yields:
+    Tuple (input_size, filter_size, out_size, stride, padding), the depthwise
+    convolution parameters.
+  """
+  input_sizes = [[4, 5, 5, 48], [2, 5, 5, 48], [4, 8, 8, 84], [4, 17, 17, 48],
+                 [4, 9, 27, 8], [4, 31, 31, 7], [4, 35, 35, 2],
+                 [4, 147, 147, 2], [3, 299, 299, 3], [5, 183, 183, 1]]
+  filter_sizes = [[1, 1, 48, 2], [2, 2, 48, 8], [1, 3, 84, 1], [3, 1, 48, 4],
+                  [3, 3, 8, 1], [3, 3, 7, 1], [5, 5, 2, 1], [3, 3, 2, 8],
+                  [2, 2, 3, 8], [5, 5, 1, 2]]
+  out_sizes = [[4, 5, 5, 96], [2, 5, 5, 384], [4, 8, 8, 84], [4, 17, 17, 192],
+               [4, 9, 27, 8], [4, 31, 31, 7], [4, 35, 35, 2], [4, 49, 49, 16],
+               [3, 150, 150, 24], [5, 92, 92, 2]]
+  strides = [1, 1, 1, 1, 1, 1, 1, 3, 2, 2]
+  # pylint: disable=invalid-name
+  VALID = "VALID"
+  SAME = "SAME"
+  # pylint: enable=invalid-name
+  paddings = [SAME, SAME, SAME, SAME, SAME, SAME, SAME, VALID, SAME, SAME, SAME]
+  for i, f, o, s, p in zip(input_sizes, filter_sizes, out_sizes, strides,
+                           paddings):
+    yield i, f, o, s, p
+
+
+def ConfigsWithDilationsToTest():
+  """Iterator for different convolution shapes, strides and paddings.
+
+  Yields:
+    Tuple (input_size, filter_size, out_size, stride, dilation, padding), the
+    depthwise
+    convolution parameters.
+  """
+  input_sizes = [[4, 6, 6, 48], [4, 8, 8, 84], [4, 36, 36, 2], [4, 148, 148, 2],
+                 [3, 300, 300, 3]]
+  filter_sizes = [[1, 1, 48, 2], [1, 3, 84, 1], [5, 5, 2, 1], [4, 4, 2, 8],
+                  [2, 2, 3, 8]]
+  out_sizes = [[4, 6, 6, 96], [4, 8, 8, 84], [4, 36, 36, 2], [4, 74, 74, 16],
+               [3, 296, 296, 24]]
+  strides = [1, 1, 2, 2, 1]
+  dilations = [2, 2, 4, 2, 4]
+  # pylint: disable=invalid-name
+  VALID = "VALID"
+  SAME = "SAME"
+  # pylint: enable=invalid-name
+  paddings = [SAME, SAME, SAME, SAME, VALID]
+  for i, f, o, s, d, p in zip(input_sizes, filter_sizes, out_sizes, strides,
+                              dilations, paddings):
+    yield i, f, o, s, d, p
+
+
+def CheckGradConfigsToTest():
+  """Iterator for different convolution shapes, strides and paddings.
+
+  compute_gradient_error() is very expensive. So the configs should be
+  relatively small.
+
+  Yields:
+    Tuple (input_size, filter_size, out_size, stride, padding), the depthwise
+    convolution parameters.
+  """
+  input_sizes = [[2, 5, 8, 1], [4, 5, 5, 1], [2, 4, 4, 2], [1, 15, 15, 2],
+                 [2, 15, 16, 1]]
+  filter_sizes = [[4, 4, 1, 2], [2, 2, 1, 2], [3, 1, 2, 2], [1, 3, 2, 1],
+                  [3, 3, 1, 2]]
+  out_sizes = [[2, 5, 8, 2], [4, 2, 2, 2], [2, 4, 4, 4], [1, 15, 15, 2],
+               [2, 5, 5, 2]]
+  strides = [1, 2, 1, 1, 3]
+  # pylint: disable=invalid-name
+  VALID = "VALID"
+  SAME = "SAME"
+  # pylint: enable=invalid-name
+  paddings = [SAME, VALID, SAME, SAME, VALID]
+  for i, f, o, s, p in zip(input_sizes, filter_sizes, out_sizes, strides,
+                           paddings):
+    yield i, f, o, s, p
+
+
+class DepthwiseConv2DTest(xla_test.XLATestCase):
+
+  # This is testing that depthwise_conv2d and depthwise_conv2d_native
+  # produce the same results.  It also tests that NCHW and NWHC
+  # formats agree, by comparing the depthwise_conv2d_native with
+  # 'NCHW' format (with transposition) matches the 'NHWC' format using
+  # the higher level interface.
+  def _VerifyValues(self,
+                    tensor_in_sizes,
+                    filter_in_sizes,
+                    stride,
+                    padding,
+                    data_type,
+                    data_format="NHWC"):
+    """Verifies the output values of the convolution function.
+
+    Args:
+      tensor_in_sizes: Input tensor dimensions in
+        [batch, input_rows, input_cols, input_depth].
+      filter_in_sizes: Filter tensor dimensions in
+        [filter_rows, filter_cols, input_depth, depth_multiplier].
+      stride: Stride.
+      padding: Padding type.
+      data_type: The data type to use.
+      data_format: The data_format of the input. "NHWC" or "NCHW".
+    """
+    total_size_1 = 1
+    total_size_2 = 1
+    for s in tensor_in_sizes:
+      total_size_1 *= s
+    for s in filter_in_sizes:
+      total_size_2 *= s
+    # Initializes the input and filter tensor with numbers incrementing from 1.
+    x1 = np.array([f * 1.0 for f in range(1, total_size_1 + 1)],
+                  dtype=data_type).reshape(tensor_in_sizes)
+    x2 = np.array([f * 1.0 for f in range(1, total_size_2 + 1)],
+                  dtype=data_type).reshape(filter_in_sizes)
+    with self.session() as sess:
+      if data_type == np.float32:
+        tolerance = 1e-4
+      else:
+        self.assertEqual(data_type, np.float64)
+        tolerance = 1e-8
+
+      t1 = array_ops.placeholder(shape=tensor_in_sizes, dtype=data_type)
+      t2 = array_ops.placeholder(shape=filter_in_sizes, dtype=data_type)
+
+      native_t1 = t1
+      strides = [1, stride, stride, 1]
+      if data_format == "NCHW":
+        # Transpose from NWHC input to NCHW
+        # Ex. [4, 5, 5, 48] to [4, 48, 5, 5]
+        native_t1 = array_ops.transpose(t1, [0, 3, 1, 2])
+        strides = [1, 1, stride, stride]
+
+      with self.test_scope():
+        conv_native = nn_ops.depthwise_conv2d_native(
+            native_t1,
+            t2,
+            strides=strides,
+            data_format=data_format,
+            padding=padding)
+
+      if data_format == "NCHW":
+        # Transpose back from NCHW to NHWC
+        conv_native = array_ops.transpose(conv_native, [0, 2, 3, 1])
+
+      with ops.device("CPU"):
+        conv_interface = ReferenceDepthwiseConv2D(
+            t1, t2, strides=[1, stride, stride, 1], padding=padding)
+
+      native_result = sess.run(conv_native, {t1: x1, t2: x2})
+      interface_result = sess.run(conv_interface, {t1: x1, t2: x2})
+
+    print("data_type:", data_type, "max diff = ",
+          np.amax(np.absolute(native_result - interface_result)))
+    self.assertAllClose(
+        np.ravel(native_result), np.ravel(interface_result), rtol=tolerance)
+
+  @test_util.run_without_tensor_float_32(
+      "DepthwiseConv2D may use TF32 when available.")
+  def testDepthwiseConv2D(self):
+    for index, (input_size, filter_size, _, stride,
+                padding) in enumerate(ConfigsToTest()):
+      print("Testing DepthwiseConv2D,", index, "th config:", input_size, "*",
+            filter_size, "stride:", stride, "padding:", padding)
+      for data_type in self.float_types:
+        # TODO(phawkins): the reference implementation only supports float32.
+        if data_type == np.float32:
+          self._VerifyValues(
+              input_size, filter_size, stride, padding, data_type)
+
+  @test_util.run_without_tensor_float_32(
+      "DepthwiseConv2D may use TF32 when available.")
+  def testDepthwiseConv2DFormat(self):
+    for index, (input_size, filter_size, _, stride,
+                padding) in enumerate(ConfigsToTest()):
+      print("Testing DepthwiseConv2DFormat,", index, "th config:", input_size,
+            "*", filter_size, "stride:", stride, "padding:", padding)
+      for data_type in self.float_types:
+        # TODO(phawkins): the reference implementation only supports float32.
+        if data_type == np.float32:
+          self._VerifyValues(
+              input_size,
+              filter_size,
+              stride,
+              padding,
+              data_type,
+              data_format="NCHW")
+
+# This is testing against hand calculated results.
+
+  def _VerifyHandValues(self, tensor_in_sizes, filter_in_sizes, stride, padding,
+                        expected):
+    """Verifies the output values of the depthwise convolution function.
+
+    Args:
+      tensor_in_sizes: Input tensor dimensions in
+        [batch, input_rows, input_cols, input_depth].
+      filter_in_sizes: Filter tensor dimensions in
+        [filter_rows, filter_cols, input_depth, depth_multiplier].
+      stride: Stride.
+      padding: Padding type.
+      expected: An array containing the expected operation outputs.
+    """
+    total_size_1 = 1
+    total_size_2 = 1
+    for s in tensor_in_sizes:
+      total_size_1 *= s
+    for s in filter_in_sizes:
+      total_size_2 *= s
+    # Initializes the input tensor with array containing incrementing
+    # numbers from 1.
+    x1 = np.array([f * 1.0 for f in range(1, total_size_1 + 1)],
+                  dtype=np.float32).reshape(tensor_in_sizes)
+    x2 = np.array([f * 1.0 for f in range(1, total_size_2 + 1)],
+                  dtype=np.float32).reshape(filter_in_sizes)
+    with self.session() as sess:
+      t1 = array_ops.placeholder(shape=tensor_in_sizes, dtype=np.float32)
+      t2 = array_ops.placeholder(shape=filter_in_sizes, dtype=np.float32)
+      with self.test_scope():
+        conv = nn_ops.depthwise_conv2d_native(
+            t1, t2, strides=[1, stride, stride, 1], padding=padding)
+      value = sess.run(conv, {t1: x1, t2: x2})
+    print("value = ", value)
+    self.assertArrayNear(expected, np.ravel(value), 1e-4)
+    self.assertShapeEqual(value, conv)
+
+  def testConv2D2x2Filter(self):
+    # The inputs look like this (it's a 3 x 2 matrix, each of depth 2):
+    #
+    # [ (1.0, 2.0), (3.0,  4.0), ( 5.0,  6.0) ]
+    # [ (7.0, 8.0), (9.0, 10.0), (11.0, 12.0) ]
+    #  We can view this as two inputs
+    #
+    #  input depth 0:
+    #
+    #  [ 1.0,  3.0,  5.0 ]
+    #  [ 7.0,  9.0, 11.0 ]
+    #
+    #  input depth 1:
+    #
+    #  [ 2.0,  4.0,  6.0 ]
+    #  [ 8.0, 10.0, 12.0 ]
+    #
+    # The filter looks like this (it has two 2 x 2 patches, each generating 2
+    # depths):
+    #
+    #  filter #0:
+    #
+    #  [ (1.0,  3.0), ( 5.0,  7.0)]
+    #  [ (9.0, 11.0), (13.0, 15.0)]
+    #
+    #  filter #1:
+    #
+    #  [ ( 2.0,  4.0), ( 6.0,  8.0)]
+    #  [ (10.0, 12.0), (14.0, 16.0)]
+    #
+    # So the outputs are:
+    #
+    # (position 0, 0: in_depth 0, output_depth 0 -- using filter #0)
+    #  1.0 * 1.0 + 7.0 * 9.0 + 3.0 * 5.0 + 9.0 * 13.0 = 196
+    # (position 0, 0: in_depth 0, output_depth 1 -- using filter #1)
+    #  1.0 * 2.0 + 7.0 * 10.0 + 3.0 * 6.0 + 9.0 * 14.0 = 216
+    # (position 0, 0: in_depth 1, output_depth 2 -- using filter #0)
+    #  2.0 * 3.0 + 8.0 * 11.0 + 4.0 * 7.0 + 10.0 * 15.0 = 272
+    # (position 0, 0: in_depth 1, output_depth 3 -- using filter #1)
+    #  2.0 * 4.0 + 8.0 * 12.0 + 4.0 * 8.0 + 10.0 * 16.0 = 296
+    #
+    # (position 1, 0: in_depth 0, output_depth 0 -- using filter #0)
+    #  3.0 * 1.0 + 9.0 * 9.0 + 5.0 * 5.0 + 11.0 * 13.0 = 252
+    # (position 1, 0: in_depth 0, output_depth 1 -- using filter #1)
+    #  3.0 * 2.0 + 9.0 * 10.0 + 5.0 * 6.0 + 11.0 * 14.0 = 280
+    # (position 1, 0: in_depth 1, output_depth 2 -- using filter #0)
+    #  4.0 * 3.0 + 10.0 * 11.0 + 6.0 * 7.0 + 12.0 * 15.0 = 344
+    # (position 1, 0: in_depth 1, output_depth 3 -- using filter #1)
+    #  4.0 * 4.0 + 10.0 * 12.0 + 6.0 * 8.0 + 12.0 * 16.0 = 376
+    expected_output = [196, 216, 272, 296, 252, 280, 344, 376]
+    self._VerifyHandValues(
+        tensor_in_sizes=[1, 2, 3, 2],
+        filter_in_sizes=[2, 2, 2, 2],
+        stride=1,
+        padding="VALID",
+        expected=expected_output)
+
+  # This is testing that depthwise_conv2d with dilation produces
+  # the same results between CPU and TPU. It also tests that NCHW
+  # and NWHC formats agree.
+  def _VerifyValuesWithDilation(self,
+                                tensor_in_sizes,
+                                filter_in_sizes,
+                                stride,
+                                dilation,
+                                padding,
+                                data_type,
+                                data_format="NHWC"):
+    """Verifies the output values of the convolution function.
+
+    Args:
+      tensor_in_sizes: Input tensor dimensions in [batch, input_rows,
+        input_cols, input_depth].
+      filter_in_sizes: Filter tensor dimensions in [filter_rows, filter_cols,
+        input_depth, depth_multiplier].
+      stride: Stride.
+      dilation: Dilation.
+      padding: Padding type.
+      data_type: The data type to use.
+      data_format: The data_format of the input. "NHWC" or "NCHW".
+    """
+    total_size_1 = 1
+    total_size_2 = 1
+    for s in tensor_in_sizes:
+      total_size_1 *= s
+    for s in filter_in_sizes:
+      total_size_2 *= s
+    # Initializes the input and filter tensor with numbers incrementing from 1.
+    x1 = np.array([f * 1.0 for f in range(1, total_size_1 + 1)],
+                  dtype=data_type).reshape(tensor_in_sizes)
+    x2 = np.array([f * 1.0 for f in range(1, total_size_2 + 1)],
+                  dtype=data_type).reshape(filter_in_sizes)
+    with self.session() as sess:
+      if data_type == np.float32:
+        # TODO(b/64210055): Tolerance for TPU is high.
+        tolerance = 1e-2
+      else:
+        self.assertEqual(data_type, np.float64)
+        tolerance = 1e-8
+
+      t1 = array_ops.placeholder(shape=tensor_in_sizes, dtype=data_type)
+      t2 = array_ops.placeholder(shape=filter_in_sizes, dtype=data_type)
+
+      native_t1 = t1
+      strides = [1, stride, stride, 1]
+      dilations = [dilation, dilation]
+      if data_format == "NCHW":
+        # Transpose from NWHC input to NCHW
+        # Ex. [4, 5, 5, 48] to [4, 48, 5, 5]
+        native_t1 = array_ops.transpose(t1, [0, 3, 1, 2])
+        strides = [1, 1, stride, stride]
+
+      with self.test_scope():
+        conv_native = nn_impl.depthwise_conv2d(
+            native_t1,
+            t2,
+            strides=strides,
+            rate=dilations,
+            data_format=data_format,
+            padding=padding)
+
+      if data_format == "NCHW":
+        # Transpose back from NCHW to NHWC
+        conv_native = array_ops.transpose(conv_native, [0, 2, 3, 1])
+
+      with ops.device("CPU"):
+        # CPU only support NHWC format
+        strides = [1, stride, stride, 1]
+        conv_interface = nn_impl.depthwise_conv2d(
+            t1, t2, strides=strides, rate=dilations, padding=padding)
+
+      native_result = sess.run(conv_native, {t1: x1, t2: x2})
+      interface_result = sess.run(conv_interface, {t1: x1, t2: x2})
+
+    print("data_type:", data_type, "max diff = ",
+          np.amax(np.absolute(native_result - interface_result)))
+    self.assertAllClose(
+        np.ravel(native_result), np.ravel(interface_result), rtol=tolerance)
+
+  def testDilationDepthwiseConv2DWith(self):
+    for index, (input_size, filter_size, _, stride, dilation,
+                padding) in enumerate(ConfigsWithDilationsToTest()):
+      print("Testing DilationDepthwiseConv2D,", index, "th config:", input_size,
+            "*", filter_size, "stride:", stride, "dilation: ", dilation,
+            "padding:", padding)
+      for data_type in self.float_types:
+        # TODO(phawkins): the reference implementation only supports float32.
+        if data_type == np.float32:
+          self._VerifyValuesWithDilation(input_size, filter_size, stride,
+                                         dilation, padding, data_type)
+
+  def testDilationDepthwiseConv2DWithFormat(self):
+    for index, (input_size, filter_size, _, stride, dilation,
+                padding) in enumerate(ConfigsWithDilationsToTest()):
+      print("Testing DilationDepthwiseConv2DFormat,", index, "th config:",
+            input_size, "*", filter_size, "stride:", stride, "dilation:",
+            dilation, "padding:", padding)
+      for data_type in self.float_types:
+        # TODO(phawkins): the reference implementation only supports float32.
+        if data_type == np.float32:
+          self._VerifyValuesWithDilation(
+              input_size,
+              filter_size,
+              stride,
+              dilation,
+              padding,
+              data_type,
+              data_format="NCHW")
+
+  def _CompareBackpropInput(self, input_sizes, filter_sizes, output_sizes,
+                            stride, padding):
+    x1 = np.random.rand(*filter_sizes).astype(np.float32)
+    x2 = np.random.rand(*output_sizes).astype(np.float32)
+
+    def _GetVal(use_xla):
+      with self.session():
+        t0 = constant_op.constant(input_sizes, shape=[len(input_sizes)])
+        t1 = array_ops.placeholder(np.float32, shape=filter_sizes)
+        t2 = array_ops.placeholder(np.float32, shape=output_sizes)
+        if use_xla:
+          with self.test_scope():
+            backprop = nn_ops.depthwise_conv2d_native_backprop_input(
+                t0, t1, t2, strides=[1, stride, stride, 1], padding=padding)
+        else:
+          backprop = nn_ops.depthwise_conv2d_native_backprop_input(
+              t0, t1, t2, strides=[1, stride, stride, 1], padding=padding)
+
+        ret = backprop.eval({t1: x1, t2: x2})
+        self.assertShapeEqual(ret, backprop)
+        return ret
+
+    gpu_value = _GetVal(use_xla=True)
+    cpu_value = _GetVal(use_xla=False)
+    self.assertAllClose(cpu_value, gpu_value, rtol=1e-3, atol=1e-3)
+
+  def testDepthwiseConv2DInputGradCompare(self):
+    for index, (input_size, filter_size, output_size, stride,
+                padding) in enumerate(ConfigsToTest()):
+      print("Testing DepthwiseConv2DInputGradCompare,", index, "th config:",
+            input_size, "*", filter_size, "stride:", stride, "padding:",
+            padding)
+      self._CompareBackpropInput(input_size, filter_size, output_size, stride,
+                                 padding)
+
+  def _CompareBackpropFilter(self,
+                             input_sizes,
+                             filter_sizes,
+                             output_sizes,
+                             stride,
+                             padding,
+                             data_format="NHWC"):
+    x0 = np.random.rand(*input_sizes).astype(np.float32)
+    x2 = np.random.rand(*output_sizes).astype(np.float32)
+
+    def _GetVal(use_xla):
+      with self.session():
+        t0 = array_ops.placeholder(np.float32, shape=input_sizes)
+        t1 = constant_op.constant(filter_sizes, shape=[len(filter_sizes)])
+        t2 = array_ops.placeholder(np.float32, shape=output_sizes)
+        native_t0 = t0
+        native_t2 = t2
+        strides = [1, stride, stride, 1]
+
+        if use_xla:
+          if data_format == "NCHW":
+            # Transpose from NWHC input to NCHW
+            # Ex. [4, 5, 5, 48] to [4, 48, 5, 5]
+            native_t0 = array_ops.transpose(t0, [0, 3, 1, 2])
+            native_t2 = array_ops.transpose(t2, [0, 3, 1, 2])
+            strides = [1, 1, stride, stride]
+          with self.test_scope():
+            backprop = nn_ops.depthwise_conv2d_native_backprop_filter(
+                native_t0,
+                t1,
+                native_t2,
+                strides=strides,
+                padding=padding,
+                data_format=data_format)
+        else:
+          # For CPU, the format NCHW is not supported. Therefore we always use
+          # NHWC here.
+          backprop = nn_ops.depthwise_conv2d_native_backprop_filter(
+              native_t0, t1, native_t2, strides=strides, padding=padding)
+        ret = backprop.eval({t0: x0, t2: x2})
+        self.assertShapeEqual(ret, backprop)
+        return ret
+
+    gpu_value = _GetVal(use_xla=True)
+    cpu_value = _GetVal(use_xla=False)
+    self.assertAllClose(cpu_value, gpu_value, rtol=1e-4, atol=1e-4)
+
+  @test_util.run_without_tensor_float_32(
+      "DepthwiseConv2DFilterGrad may use TF32 when available.")
+  def testDepthwiseConv2DFilterGradCompare(self):
+    for index, (input_size, filter_size, output_size, stride,
+                padding) in enumerate(ConfigsToTest()):
+      print("Testing DepthwiseConv2DFilterGradCompare,", index, "th config:",
+            input_size, "*", filter_size, "producing output", output_size,
+            "stride:", stride, "padding:", padding)
+      self._CompareBackpropFilter(input_size, filter_size, output_size,
+                                  stride, padding)
+
+  @test_util.run_without_tensor_float_32(
+      "DepthwiseConv2DFilterGrad may use TF32 when available.")
+  def testDepthwiseConv2DFilterGradFormatNCHWCompare(self):
+    for index, (input_size, filter_size, output_size, stride,
+                padding) in enumerate(ConfigsToTest()):
+      print("Testing DepthwiseConv2DFilterGradFormatNCHWCompare,", index,
+            "th config:", input_size, "*", filter_size, "producing output",
+            output_size, "stride:", stride, "padding:", padding)
+      self._CompareBackpropFilter(
+          input_size,
+          filter_size,
+          output_size,
+          stride,
+          padding,
+          data_format="NCHW")
+
+  def _CompareBackpropInputWithDilation(self, input_sizes, filter_sizes,
+                                        output_sizes, stride, dilation,
+                                        padding):
+    x1 = np.random.rand(*filter_sizes).astype(np.float32)
+    x2 = np.random.rand(*output_sizes).astype(np.float32)
+
+    def _GetVal(use_xla):
+      with self.session():
+        t1 = array_ops.placeholder(np.float32, shape=filter_sizes)
+        t2 = array_ops.placeholder(np.float32, shape=output_sizes)
+        if use_xla:
+          with self.test_scope():
+            t0 = constant_op.constant(input_sizes, shape=[len(input_sizes)])
+            backprop = nn_ops.depthwise_conv2d_native_backprop_input(
+                t0,
+                t1,
+                t2,
+                strides=[1, stride, stride, 1],
+                dilations=[1, dilation, dilation, 1],
+                padding=padding)
+        else:
+          # TODO(wangtao): figure out gradient with stride > 1.
+          # depthwise_conv2d_native_backprop_input on CPU doesn't support
+          # dilation.
+          t3 = array_ops.space_to_batch(
+              t2, block_size=dilation, paddings=[[0, 0], [0, 0]])
+          input_sizes_transform = [
+              input_sizes[0] * dilation * dilation, input_sizes[1] // dilation,
+              input_sizes[2] // dilation, input_sizes[3]
+          ]
+          t0 = constant_op.constant(
+              input_sizes_transform, shape=[len(input_sizes)])
+          backprop_naive = nn_ops.depthwise_conv2d_native_backprop_input(
+              t0, t1, t3, strides=[1, stride, stride, 1], padding=padding)
+          backprop = array_ops.batch_to_space(
+              backprop_naive, [[0, 0], [0, 0]], block_size=dilation)
+
+        ret = backprop.eval({t1: x1, t2: x2})
+        self.assertShapeEqual(ret, backprop)
+        return ret
+
+    gpu_value = _GetVal(use_xla=True)
+    cpu_value = _GetVal(use_xla=False)
+
+    # TODO (b/64210055): Tolerance for TPU is high.
+    self.assertAllClose(cpu_value, gpu_value, rtol=1e-2, atol=1e-3)
+
+  def testDilationDepthwiseConv2DInputGradWithCompare(self):
+    for index, (input_size, filter_size, output_size, stride, dilation,
+                padding) in enumerate(ConfigsWithDilationsToTest()):
+      print("Testing DilationDepthwiseConv2DInputGradWithDilationCompare,",
+            index, "th config:", input_size, "*", filter_size, "stride:",
+            stride, "dilation:", dilation, "padding:", padding)
+      # TODO(wangtao): implement CPU grad computation with stride > 1.
+      if stride == 1:
+        self._CompareBackpropInputWithDilation(input_size, filter_size,
+                                               output_size, stride, dilation,
+                                               padding)
+
+  def _CompareBackpropFilterWithDilation(self,
+                                         input_sizes,
+                                         filter_sizes,
+                                         output_sizes,
+                                         stride,
+                                         dilation,
+                                         padding,
+                                         data_format="NHWC"):
+    x0 = np.random.rand(*input_sizes).astype(np.float32)
+    x2 = np.random.rand(*output_sizes).astype(np.float32)
+
+    def _GetVal(use_xla):
+      with self.session():
+        t0 = array_ops.placeholder(np.float32, shape=input_sizes)
+        t1 = constant_op.constant(filter_sizes, shape=[len(filter_sizes)])
+        t2 = array_ops.placeholder(np.float32, shape=output_sizes)
+        native_t0 = t0
+        native_t2 = t2
+        strides = [1, stride, stride, 1]
+        dilations = [1, dilation, dilation, 1]
+
+        if use_xla:
+          if data_format == "NCHW":
+            # Transpose from NWHC input to NCHW
+            # Ex. [4, 5, 5, 48] to [4, 48, 5, 5]
+            native_t0 = array_ops.transpose(t0, [0, 3, 1, 2])
+            native_t2 = array_ops.transpose(t2, [0, 3, 1, 2])
+            strides = [1, 1, stride, stride]
+            dilations = [1, 1, dilation, dilation]
+          with self.test_scope():
+            backprop = nn_ops.depthwise_conv2d_native_backprop_filter(
+                native_t0,
+                t1,
+                native_t2,
+                strides=strides,
+                padding=padding,
+                dilations=dilations,
+                data_format=data_format)
+        else:
+          # For CPU, the format NCHW is not supported. Therefore we always use
+          # NHWC here.
+          # depthwise_conv2d_native_backprop_filter on CPU doesn't support
+          # dilation.
+          native_t3 = array_ops.space_to_batch(
+              native_t2, block_size=dilation, paddings=[[0, 0], [0, 0]])
+          native_t0_transform = array_ops.space_to_batch(
+              native_t0, block_size=dilation, paddings=[[0, 0], [0, 0]])
+          backprop = nn_ops.depthwise_conv2d_native_backprop_filter(
+              native_t0_transform,
+              t1,
+              native_t3,
+              strides=strides,
+              padding=padding)
+        ret = backprop.eval({t0: x0, t2: x2})
+        self.assertShapeEqual(ret, backprop)
+        return ret
+
+    gpu_value = _GetVal(use_xla=True)
+    cpu_value = _GetVal(use_xla=False)
+    # TODO(b/64210055): Tolerance for TPU is high.
+    self.assertAllClose(cpu_value, gpu_value, rtol=1e-3, atol=1e-4)
+
+  def testDilationDepthwiseConv2DFilterGradCompare(self):
+    for index, (input_size, filter_size, output_size, stride, dilation,
+                padding) in enumerate(ConfigsWithDilationsToTest()):
+      print("Testing DilationDepthwiseConv2DFilterGradCompare,", index,
+            "th config:", input_size, "*", filter_size, "producing output",
+            output_size, "stride:", stride, "dilation:", dilation, "padding:",
+            padding)
+      if stride == 1:
+        # TODO(wangtao): implement CPU grad computation with stride > 1.
+        self._CompareBackpropFilterWithDilation(input_size, filter_size,
+                                                output_size, stride, dilation,
+                                                padding)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/dynamic_slice_ops_test.py

@@ -0,0 +1,89 @@
+# Copyright 2018 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 dynamic slicing ops."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.compiler.tf2xla.python import xla
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class DynamicUpdateSliceOpsTest(xla_test.XLATestCase):
+
+  def _assertOpOutputMatchesExpected(self, op, args, expected):
+    with self.session() as session:
+      with self.test_scope():
+        placeholders = [
+            array_ops.placeholder(dtypes.as_dtype(arg.dtype), arg.shape)
+            for arg in args
+        ]
+        feeds = {placeholders[i]: args[i] for i in range(0, len(args))}
+        output = op(*placeholders)
+      result = session.run(output, feeds)
+      self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testUpdateSlice(self):
+    for dtype in self.numeric_types:
+      self._assertOpOutputMatchesExpected(
+          xla.dynamic_update_slice, [
+              np.array([], dtype=dtype),
+              np.array([], dtype=dtype),
+              np.array([0], dtype=np.int32)
+          ],
+          expected=np.array([], dtype=dtype))
+
+      self._assertOpOutputMatchesExpected(
+          xla.dynamic_update_slice, [
+              np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=dtype),
+              np.array([11, 12, 13], dtype=dtype),
+              np.array([6], dtype=np.int32)
+          ],
+          expected=np.array([1, 2, 3, 4, 5, 6, 11, 12, 13, 10], dtype=dtype))
+
+      self._assertOpOutputMatchesExpected(
+          xla.dynamic_update_slice, [
+              np.array(
+                  [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype),
+              np.array([[42, 43], [44, 45]], dtype=dtype),
+              np.array([1, 2], dtype=np.int32)
+          ],
+          expected=np.array(
+              [[1, 2, 3, 4], [5, 6, 42, 43], [9, 10, 44, 45]], dtype=dtype))
+
+      self._assertOpOutputMatchesExpected(
+          xla.dynamic_update_slice, [
+              np.array(
+                  [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype),
+              np.array([[], []], dtype=dtype),
+              np.array([1, 2], dtype=np.int32)
+          ],
+          expected=np.array(
+              [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype))
+
+      self._assertOpOutputMatchesExpected(
+          xla.dynamic_update_slice, [
+              np.array(
+                  [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]], dtype=dtype),
+              np.ones([3, 4], dtype=dtype),
+              np.array([0, 0], dtype=np.int32)
+          ],
+          expected=np.ones([3, 4], dtype=dtype))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/tests/dynamic_stitch_test.py

@@ -0,0 +1,91 @@
+# 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.
+# ==============================================================================
+"""Tests for tf.dynamic_stitch."""
+
+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 data_flow_ops
+from tensorflow.python.platform import googletest
+
+
+class DynamicStitchTest(xla_test.XLATestCase):
+
+  def _AssertDynamicStitchResultIs(self, indices, data, expected):
+    with self.session() as session:
+      index_placeholders = [
+          array_ops.placeholder(dtypes.as_dtype(arg.dtype)) for arg in indices
+      ]
+      data_placeholders = [
+          array_ops.placeholder(dtypes.as_dtype(arg.dtype)) for arg in data
+      ]
+      with self.test_scope():
+        output = data_flow_ops.dynamic_stitch(index_placeholders,
+                                              data_placeholders)
+
+      feed_dict = {}
+      for placeholder, value in zip(index_placeholders, indices):
+        feed_dict[placeholder] = value
+      for placeholder, value in zip(data_placeholders, data):
+        feed_dict[placeholder] = value
+      result = session.run(output, feed_dict=feed_dict)
+      self.assertAllClose(expected, result, rtol=1e-3)
+
+  def testSimpleEmpty(self):
+    idx1 = np.array([0, 2], dtype=np.int32)
+    idx2 = np.array([[1], [3]], dtype=np.int32)
+    val1 = np.array([[], []], dtype=np.int32)
+    val2 = np.array([[[]], [[]]], dtype=np.int32)
+    self._AssertDynamicStitchResultIs(
+        [idx1, idx2], [val1, val2],
+        expected=np.array([[], [], [], []], np.int32))
+
+  def testEmptyIndex(self):
+    idx1 = np.array([], dtype=np.int32)
+    idx2 = np.array([[], []], dtype=np.int32)
+    val1 = np.ndarray(shape=(0, 9), dtype=np.int32)
+    val2 = np.ndarray(shape=(2, 0, 9), dtype=np.int32)
+    self._AssertDynamicStitchResultIs([idx1, idx2], [val1, val2],
+                                      expected=np.ndarray(
+                                          shape=(0, 9), dtype=np.int32))
+
+  def testSimple1D(self):
+    val1 = np.array([0, 4, 7], dtype=np.int32)
+    val2 = np.array([1, 6, 2, 3, 5], dtype=np.int32)
+    val3 = np.array([0, 40, 70], dtype=np.float32)
+    val4 = np.array([10, 60, 20, 30, 50], dtype=np.float32)
+    expected = np.array([0, 10, 20, 30, 40, 50, 60, 70], dtype=np.float32)
+    self._AssertDynamicStitchResultIs(
+        [val1, val2], [val3, val4], expected=expected)
+
+  def testSimple2D(self):
+    val1 = np.array([0, 4, 7], dtype=np.int32)
+    val2 = np.array([1, 6], dtype=np.int32)
+    val3 = np.array([2, 3, 5], dtype=np.int32)
+    val4 = np.array([[0, 1], [40, 41], [70, 71]], dtype=np.float32)
+    val5 = np.array([[10, 11], [60, 61]], dtype=np.float32)
+    val6 = np.array([[20, 21], [30, 31], [50, 51]], dtype=np.float32)
+    expected = np.array(
+        [[0, 1], [10, 11], [20, 21], [30, 31], [40, 41], [50, 51], [60, 61],
+         [70, 71]],
+        dtype=np.float32)
+    self._AssertDynamicStitchResultIs(
+        [val1, val2, val3], [val4, val5, val6], expected=expected)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/eager_test.py

@@ -0,0 +1,803 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Test cases for eager execution using XLA."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+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 indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.layers import convolutional
+from tensorflow.python.layers import pooling
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack  # pylint: disable=g-direct-tensorflow-import
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import functional_ops
+from tensorflow.python.ops import gen_random_ops
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import googletest
+from tensorflow.python.training import adam
+
+
+class EagerTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    with self.test_scope():
+      three = constant_op.constant(3)
+      five = constant_op.constant(5)
+      product = three * five
+      self.assertAllEqual(15, product)
+
+  def testGradientTape(self):
+    with self.test_scope():
+
+      x = constant_op.constant(1.0)
+      y = constant_op.constant(10.0)
+      with backprop.GradientTape(persistent=True) as tape:
+        tape.watch(x)
+        tape.watch(y)
+        a = x + y + x * y
+      da_dx = tape.gradient(a, x)
+      da_dy = tape.gradient(a, y)
+
+    self.assertEqual(11.0, da_dx.numpy())
+    self.assertEqual(2.0, da_dy.numpy())
+
+  def testExecuteListOutputLen0(self):
+    with self.test_scope():
+      empty = constant_op.constant([], dtype=dtypes.float32)
+      result = array_ops_stack.unstack(empty, 0)
+      self.assertTrue(isinstance(result, list))
+      self.assertEqual(0, len(result))
+
+  def testExecuteListOutputLen1(self):
+    with self.test_scope():
+      split_dim = constant_op.constant(1)
+      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
+      result = array_ops.split(value, 1, axis=split_dim)
+      self.assertTrue(isinstance(result, list))
+      self.assertEqual(1, len(result))
+      self.assertAllEqual([[0, 1, 2], [3, 4, 5]], result[0])
+
+  def testExecuteListOutputLen3(self):
+    with self.test_scope():
+      split_dim = constant_op.constant(1)
+      value = constant_op.constant([[0., 1., 2.], [3., 4., 5.]])
+      result = array_ops.split(value, 3, axis=split_dim)
+      self.assertTrue(isinstance(result, list))
+      self.assertEqual(3, len(result))
+      self.assertAllEqual([[0], [3]], result[0])
+      self.assertAllEqual([[1], [4]], result[1])
+      self.assertAllEqual([[2], [5]], result[2])
+
+  def testBasicGraph(self):
+    # Run some ops eagerly
+    with self.test_scope():
+      three = constant_op.constant(3)
+      five = constant_op.constant(5)
+      product = three * five
+      self.assertAllEqual(15, product)
+
+    # Run some ops graphly
+    with context.graph_mode(), self.session():
+      with self.test_scope():
+        three = constant_op.constant(3)
+        five = constant_op.constant(5)
+        product = three * five
+        self.assertAllEqual(15, self.evaluate(product))
+
+  def testDegenerateSlices(self):
+    with self.test_scope():
+      npt = np.arange(1, 19, dtype=np.float32).reshape(3, 2, 3)
+      t = constant_op.constant(npt)
+      # degenerate by offering a forward interval with a negative stride
+      self.assertAllEqual(npt[0:-1:-1, :, :], t[0:-1:-1, :, :])
+      # degenerate with a reverse interval with a positive stride
+      self.assertAllEqual(npt[-1:0, :, :], t[-1:0, :, :])
+      # empty interval in every dimension
+      self.assertAllEqual(npt[-1:0, 2:2, 2:3:-1], t[-1:0, 2:2, 2:3:-1])
+
+  def testIdentity(self):
+    with self.test_scope():
+      self.assertAllEqual(2, array_ops.identity(2))
+
+  def testRandomOps(self):
+    with self.test_scope():
+      tensor = gen_random_ops.random_uniform((2, 2), dtypes.float32)
+      row0 = tensor[0].numpy()
+      row1 = tensor[1].numpy()
+      # It should be very unlikely to rng to generate two equal rows.
+      self.assertFalse((row0 == row1).all())
+
+  def testIdentityOnVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(True)
+      i = array_ops.identity(v)
+    self.assertAllEqual(True, i.numpy())
+
+  def testAssignAddVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      v.assign_add(2.0)
+    self.assertEqual(3.0, v.numpy())
+
+  def testReadAssignRead(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      val1 = v.read_value()
+      v.assign_add(2.0)
+      val2 = v.read_value()
+    self.assertEqual(1.0, val1.numpy())
+    self.assertEqual(3.0, val2.numpy())
+
+  def testGradient(self):
+    def f(x):
+      return x
+
+    with self.test_scope():
+      grad_fn = backprop.gradients_function(f)
+      self.assertAllEqual(2., grad_fn(1., dy=2.)[0])
+
+  def testVariableGradient(self):
+    with self.test_scope():
+      v0 = resource_variable_ops.ResourceVariable(1.0)
+
+      def f():
+        x = v0 * v0
+        return x
+
+      grads = backprop.implicit_grad(f)()
+    self.assertEqual(2., grads[0][0].numpy())
+
+  def testMultipleVariableReads(self):
+    # This test makes sure consecutive variable reads don't copy
+    # the underlying memory.
+    with self.test_scope():
+      # Create 128MiB variables
+      var = resource_variable_ops.ResourceVariable(
+          array_ops.ones([32, 1024, 1024]))
+
+      # Read the same variable 100 times. If the underlying tensor
+      # is not copied, this is a trivial operation. If it is copied,
+      # this will eat over 13GB and OOM.
+      values = []
+      for _ in range(100):
+        values.append(var.value())
+
+  # The shape, shape_n, size, and rank are tested here because their
+  # execution kernels (as opposed to compilation only tf2xla kernels)
+  # are distincts from tf2xla kernels.
+
+  def testShape(self):
+    def const(value):
+      return array_ops.shape(
+          constant_op.constant(value)).numpy()
+
+    def ones(value):
+      return array_ops.shape(
+          array_ops.ones(value)).numpy()
+
+    with self.test_scope():
+      # Shapes of directly constructed tensors
+      self.assertAllEqual([], const(3))
+      self.assertAllEqual([3], const([1.0, 2.0, 3.0]))
+      self.assertAllEqual([2, 2], const([[1.0, 2.0], [3.0, 4.0]]))
+      self.assertAllEqual([2, 1, 2], const([[[1.0, 2.0]], [[3.0, 4.0]]]))
+
+      # Shapes of tensors created by op running on device
+      # We make this distinction because directly constructed tensors
+      # are treated differently in a few places that can influence shape:
+      #  - they always have on_host_tensor
+      #  - they and their shapes can be cached
+      #  - they end up on device via a copy, instead of as program output
+      self.assertAllEqual([], ones([]))
+      self.assertAllEqual([3], ones([3]))
+      self.assertAllEqual([2, 2], ones([2, 2]))
+      self.assertAllEqual([2, 1, 2], ones([2, 1, 2]))
+
+  def testShapeN(self):
+    with self.test_scope():
+      # Shapes of directly constructed tensors
+      shapes = array_ops.shape_n([
+          constant_op.constant(1.0),
+          constant_op.constant([1.0, 2.0, 3.0]),
+          constant_op.constant([[1.0, 2.0], [3.0, 4.0]])])
+      self.assertAllEqual(
+          [[], [3], [2, 2]],
+          [x.numpy().tolist() for x in shapes])
+
+      # Shapes of tensors created by op running on device
+      shapes = array_ops.shape_n([
+          array_ops.ones([]),
+          array_ops.ones([3]),
+          array_ops.ones([2, 2])])
+      self.assertAllEqual(
+          [[], [3], [2, 2]],
+          [x.numpy().tolist() for x in shapes])
+
+  def testSize(self):
+    with self.test_scope():
+      self.assertEqual(
+          1, array_ops.size(constant_op.constant(1.0)).numpy())
+      self.assertEqual(
+          3, array_ops.size(constant_op.constant([1.0, 2.0, 3.0])).numpy())
+      self.assertEqual(
+          4, array_ops.size(
+              constant_op.constant([[1.0, 2.0], [3.0, 4.0]])).numpy())
+
+  def testRank(self):
+    with self.test_scope():
+      self.assertEqual(
+          0, array_ops.rank(constant_op.constant(1.0)).numpy())
+      self.assertEqual(
+          1, array_ops.rank(constant_op.constant([1.0, 2.0, 3.0])).numpy())
+      self.assertEqual(
+          2, array_ops.rank(
+              constant_op.constant([[1.0, 2.0], [3.0, 4.0]])).numpy())
+
+  def testAdam(self):
+    with self.test_scope():
+      optimizer = adam.AdamOptimizer(0.1)
+      x = resource_variable_ops.ResourceVariable(10.0)
+      with backprop.GradientTape() as tape:
+        y = x * x
+      dy_dx = tape.gradient(y, x)
+      optimizer.apply_gradients([(dy_dx, x)])
+      self.assertAlmostEqual(9.9, x.numpy(), places=3)
+
+  def testAdamSparse(self):
+    with ops.device('/cpu:0'):
+      # Create 2-D embedding for 3 objects on CPU because sparse/sliced updates
+      # are not implemented on TPU.
+      embedding_matrix = resource_variable_ops.ResourceVariable(
+          array_ops.ones([3, 2]))
+
+    with self.test_scope():
+      with backprop.GradientTape() as tape:
+        embedding = embedding_ops.embedding_lookup(embedding_matrix, [1])
+        y = math_ops.reduce_sum(embedding)
+      dy_dx = tape.gradient(y, embedding_matrix)
+      self.assertIsInstance(dy_dx, indexed_slices.IndexedSlices)
+      optimizer = adam.AdamOptimizer(0.1)
+      # The gradient application operations will run on CPU because optimizer
+      # updates are always collocated with the variable.
+      optimizer.apply_gradients([(dy_dx, embedding_matrix)])
+
+      # This assign_add will run on CPU because when an input to an
+      # operation is a resource, this operation is placed on the resource's
+      # device by the eager runtime.
+      embedding_matrix.assign_add(array_ops.ones([3, 2]))
+
+    self.assertAllClose([[2.0, 2.0],
+                         [1.9, 1.9],
+                         [2.0, 2.0]], embedding_matrix.numpy())
+
+
+class EagerFunctionTest(xla_test.XLATestCase):
+
+  def testBasic(self):
+    with self.test_scope():
+      matmul = def_function.function(math_ops.matmul)
+      t = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+      sq = matmul(t, t, transpose_a=True)
+      self.assertAllEqual(sq.numpy().reshape(-1), [10, 14, 14, 20])
+
+  def testConv(self):
+    if 'GPU' in self.device:
+      # TODO(b/32333178)
+      self.skipTest('Current implementation of RandomStandardNormal kernel '
+                    'is very slow on GPU, and has been denylisted.')
+    with self.test_scope():
+      data_format = 'channels_last'
+      conv = convolutional.Conv2D(
+          filters=1, kernel_size=2, padding='VALID',
+          data_format=data_format, activation=nn_ops.relu,
+          kernel_initializer=init_ops.ones_initializer(),
+          bias_initializer=init_ops.zeros_initializer())
+      pool = pooling.MaxPooling2D(2, 2, data_format=data_format)
+
+      def model(x):
+        x = conv(x)
+        return pool(x)
+      model = def_function.function(model)
+
+      x = array_ops.ones([1, 4, 4, 1])
+      y = model(x)
+      self.assertAllEqual(y.numpy(), [[[[4.]]]])
+
+  def testReadVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      @def_function.function
+      def f():
+        return v.read_value()
+
+      var = f()
+      self.assertEqual(1.0, var.numpy())
+
+  def testResourceVariableNoInlineReadWrite(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+      w = resource_variable_ops.ResourceVariable(0.0)
+
+      @def_function.function(experimental_attributes={'_noinline': True})
+      def g(x):
+        w.assign(w.read_value() + x)
+        return v.read_value() + x * w.read_value()
+
+      @def_function.function(experimental_attributes={'_noinline': True})
+      def f():
+        return g(1.0) + g(2.0) + g(3.0) + g(4.0) + g(5.0)
+
+      # 1 + 1*1 + 1 + 2*3 + 1 + 3*6 + 1 + 4*10 + 1 + 5*15
+      self.assertEqual(145.0, f().numpy())
+      self.assertEqual(15.0, w.read_value().numpy())
+
+  def testResourceVariableNoInlineReadOnly(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(10.0)
+
+      @def_function.function(experimental_attributes={'_noinline': True})
+      def g():
+        return v.read_value()
+
+      @def_function.function(experimental_attributes={'_noinline': True})
+      def f():
+        return g() + g() + g() + g() + g()
+
+      self.assertEqual(50.0, f().numpy())
+
+  def testResourceVariableNoInlineWriteOnly(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(0.0)
+
+      @def_function.function(experimental_attributes={'_noinline': True})
+      def g(x):
+        v.assign(x)
+
+      @def_function.function(experimental_attributes={'_noinline': True})
+      def f():
+        g(1.0)
+        g(2.0)
+        g(3.0)
+        g(4.0)
+        g(5.0)
+
+      f()
+      self.assertEqual(5.0, v.read_value().numpy())
+
+  def testUpdateVariable(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable(1.0)
+
+      def f(v):
+        v.assign_add(1.0)
+        return v
+
+      f = def_function.function(f)
+
+      var = f(v)
+      self.assertEqual(2.0, var.numpy())
+
+  def testReturnResourceHandle(self):
+    with self.test_scope():
+      v = resource_variable_ops.ResourceVariable([[1.0, 2.0], [3.0, 4.0]])
+
+      def f(v):
+        return v.handle
+
+      f = def_function.function(f)
+      handle = f(v)
+      self.assertAllEqual(v.numpy(),
+                          resource_variable_ops.read_variable_op(
+                              handle, dtypes.float32).numpy())
+
+  def testReturnMultipleResourceHandles(self):
+    with self.test_scope():
+      v1 = resource_variable_ops.ResourceVariable(1.25)
+      v2 = resource_variable_ops.ResourceVariable(2.0)
+
+      def f(v):
+        return v.handle, 3.0 * v, v2.handle, v + v2
+
+      f = def_function.function(f)
+      v1_handle, v1_times_3, v2_handle, variable_sum = f(v1)
+      self.assertAllEqual(v1.numpy(),
+                          resource_variable_ops.read_variable_op(
+                              v1_handle, dtypes.float32).numpy())
+      self.assertEqual(3.75, v1_times_3.numpy())
+      self.assertAllEqual(v2.numpy(),
+                          resource_variable_ops.read_variable_op(
+                              v2_handle, dtypes.float32).numpy())
+      self.assertEqual(3.25, variable_sum.numpy())
+
+  def testAllArgumentKinds(self):
+    """Test a complex function that takes different argument kinds.
+
+    tf2xla machinery that translates, compiles, and runs defuns
+    classifies arguments into: compile-time constants, regular tensors,
+    and resources. This test creates a function with a mix of all these
+    kinds. Moreover, the order of function arguments is intentionally mixed up.
+
+    This also tests the case when the same argument is a compile-time constant
+    as well as used in an operation that normally expects its inputs to be
+    in device memory - addition in this case.
+    """
+    with self.test_scope():
+      def foo(c1, r1, v1, c2, v2, r2):
+        # c1 and c2 are compile-time constants
+        # r1 and r2 are regular tensors
+        # v1 and v2 are resource variables
+        a = c1 + r1
+        b = math_ops.cast(c2, dtypes.float32) + v2
+        c = array_ops.slice(v1, c1, c2)
+        d = r2 * v2
+        return a, b, c, d
+
+      foo = def_function.function(foo)
+
+      c1 = [0, 0]
+      c2 = array_ops.ones([2], dtype=dtypes.int32)
+
+      r1 = array_ops.ones([2])
+      r2 = [[2., 2.], [3., 3.]]
+
+      v1 = resource_variable_ops.ResourceVariable([[1., 2.], [3., 4.]])
+      v2 = resource_variable_ops.ResourceVariable([[10., 20.], [30., 40.]])
+
+      a, b, c, d = foo(c1, r1, v1, c2, v2, r2)
+
+      self.assertAllEqual([1, 1], a.numpy())
+      self.assertAllEqual([[11., 21.], [31., 41.]], b.numpy())
+      self.assertAllEqual([[1.]], c.numpy())
+      self.assertAllEqual([[20., 40.], [90., 120.]], d.numpy())
+
+  def testDefunInGradientTape(self):
+    with self.test_scope():
+      v0 = resource_variable_ops.ResourceVariable(5.0)
+
+      @def_function.function
+      def f(x):
+        x = v0 * v0 * x
+        return x
+
+      x = constant_op.constant(3.0)
+      with backprop.GradientTape() as tape:
+        y = f(x)
+      dy = tape.gradient(y, v0)
+
+    self.assertEqual(75, y.numpy())
+    self.assertEqual(30, dy.numpy())
+
+  def testGradientTapeInDefun(self):
+    with self.test_scope():
+      v0 = resource_variable_ops.ResourceVariable(5.0)
+
+      @def_function.function
+      def f():
+        x = constant_op.constant(1.0)
+        with backprop.GradientTape() as tape:
+          y = v0 * x
+        dy = tape.gradient(y, v0)
+        return dy
+
+      dy = f()
+      self.assertEqual(1.0, dy.numpy())
+
+  def testSliceInDefun(self):
+    with self.test_scope():
+
+      @def_function.function
+      def f(x, y):
+        return x[0::2, y:, ...]
+
+      x = array_ops.ones([2, 3, 4], dtype=dtypes.float32)
+      y = array_ops.ones([], dtype=dtypes.int32)
+      with backprop.GradientTape() as tape:
+        tape.watch(x)
+        tape.watch(y)
+        z = f(x, y)
+      dz = tape.gradient(z, x)
+
+      self.assertAllEqual(np.ones([1, 2, 4]), z.numpy())
+      self.assertAllEqual((2, 3, 4), dz.shape.as_list())
+
+  def testNestedDefun(self):
+    with self.test_scope():
+
+      @def_function.function
+      def times_two(x):
+        return 2. * x
+
+      @def_function.function
+      def two_x_plus_1(x):
+        return times_two(x) + 1.
+
+      x = constant_op.constant([2., 3., 4.])
+      y = two_x_plus_1(x)
+      self.assertAllEqual([5., 7., 9.], y.numpy())
+
+  def testNestedDefunWithVariable(self):
+    with self.test_scope():
+      v0 = resource_variable_ops.ResourceVariable(5.0)
+
+      @def_function.function
+      def g(x):
+        x = v0 * x
+        return x
+
+      @def_function.function
+      def f(x):
+        x = g(v0 * x)
+        return x
+
+      x = constant_op.constant(3.0)
+      y = f(x)
+
+    self.assertEqual(75.0, y.numpy())
+
+  def testNestedDefunInGradientTape(self):
+    with self.test_scope():
+      v0 = resource_variable_ops.ResourceVariable(5.0)
+
+      @def_function.function
+      def g(x):
+        x = v0 * x
+        return x
+
+      @def_function.function
+      def f(x):
+        x = g(v0 * x)
+        return x
+
+      x = constant_op.constant(3.0)
+      with backprop.GradientTape() as tape:
+        y = f(x)
+      dy = tape.gradient(y, v0)
+
+    self.assertEqual(75, y.numpy())
+    self.assertEqual(30, dy.numpy())
+
+  def testNestedDefunInGradientTapeDifferentVars(self):
+    with self.test_scope():
+      v0 = resource_variable_ops.ResourceVariable(5.0)
+      v1 = resource_variable_ops.ResourceVariable(3.0)
+
+      @def_function.function
+      def g(x):
+        x = v1 * x
+        return x
+
+      @def_function.function
+      def f(x):
+        x = g(v0 * x)
+        return x
+
+      x = constant_op.constant(3.0)
+      with backprop.GradientTape(persistent=True) as tape:
+        y = f(x)
+      dy_v0 = tape.gradient(y, v0)
+      dy_v1 = tape.gradient(y, v1)
+
+    self.assertEqual(45, y.numpy())
+    self.assertEqual(9, dy_v0.numpy())
+    self.assertEqual(15, dy_v1.numpy())
+
+  def testWhileInDefun(self):
+    with self.test_scope():
+      @def_function.function
+      def f(start):
+        c = lambda x: math_ops.less(x, 13.0)
+        b = lambda x: math_ops.add(x, 1.0)
+        return while_loop.while_loop(c, b, [start])
+
+      y = f(constant_op.constant(3.0))
+    self.assertEqual(13.0, y.numpy())
+
+  def testAutoGraphWhileInDefun(self):
+    with self.test_scope():
+      @def_function.function
+      def f(start):
+        x = start
+        while x < 13.0:
+          x += 1.0
+        return x
+
+      y = f(constant_op.constant(3.0))
+    self.assertEqual(13.0, y.numpy())
+
+  def testCondInDefun(self):
+    with self.test_scope():
+      @def_function.function
+      def f(pred, value):
+        fn1 = lambda: math_ops.add(value, 1.0)
+        fn2 = lambda: math_ops.subtract(value, 1.0)
+        return cond.cond(pred, fn1, fn2)
+
+      plus_one = f(constant_op.constant(True), constant_op.constant(10.0))
+      minus_one = f(constant_op.constant(False), constant_op.constant(10.0))
+    self.assertEqual(11.0, plus_one.numpy())
+    self.assertEqual(9.0, minus_one.numpy())
+
+  def testAutoGraphCondInDefun(self):
+    with self.test_scope():
+      @def_function.function
+      def f(pred, value):
+        if pred:
+          return value + 1.0
+        else:
+          return value - 1.0
+
+      plus_one = f(constant_op.constant(True), constant_op.constant(10.0))
+      minus_one = f(constant_op.constant(False), constant_op.constant(10.0))
+    self.assertEqual(11.0, plus_one.numpy())
+    self.assertEqual(9.0, minus_one.numpy())
+
+  def testScanInDefun(self):
+    with self.test_scope():
+      elems = constant_op.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name='data')
+      v = constant_op.constant(2.0, name='v')
+
+      @def_function.function
+      def f(y):
+        # pylint: disable=unnecessary-lambda
+        return functional_ops.scan(
+            lambda a, x: math_ops.multiply(a, x), y, initializer=v)
+        # pylint: enable=unnecessary-lambda
+
+      r = f(elems)
+      self.assertAllEqual([2., 4., 12., 48., 240., 1440.], self.evaluate(r))
+
+  def testFeedDeviceMemoryToOpExpectingHostMemory(self):
+    @def_function.function
+    def f(dims, value):
+      return array_ops.fill(dims, value)
+
+    with self.test_scope():
+      x = constant_op.constant([4], dtype=dtypes.int64)
+
+    y = f(x, 3)
+    self.assertAllEqual([3, 3, 3, 3], y)
+
+  def testRequestNotToCompile(self):
+    with self.test_scope():
+      def f(x):
+        with ops.device('device:CPU:0'):
+          y = 2.0 * x
+        return x, y
+
+      wholly_compiled_f = def_function.function(f)
+      op_by_op_f = def_function.function(f, jit_compile=False)
+
+      x = array_ops.identity([0.0, 2.0], name='data')
+
+      # When function is wholly compiled, all outputs will be on the
+      # device on which it is run.
+      r_x, r_y = wholly_compiled_f(x)
+      self.assertAllEqual([0.0, 2.0], r_x)
+      self.assertAllEqual([0.0, 4.0], r_y)
+      if context.executing_eagerly():
+        # backing_device is only available for eager tensors.
+        self.assertRegex(r_x.backing_device, self.device)
+        self.assertRegex(r_y.backing_device, self.device)
+
+      # When function is executed op-by-op, requested devices will be
+      # respected.
+      r_x, r_y = op_by_op_f(x)
+      self.assertAllEqual([0.0, 2.0], r_x)
+      self.assertAllEqual([0.0, 4.0], r_y)
+      if context.executing_eagerly():
+        # backing_device is only available for eager tensors.
+        self.assertRegex(r_x.backing_device, self.device)
+        self.assertRegex(r_y.backing_device, 'device:CPU:0')
+
+
+class ExcessivePaddingTest(xla_test.XLATestCase):
+  """Test that eager execution works with TPU flattened tensors.
+
+  Tensors that would normally be excessively padded when written
+  to TPU memory are reshaped to 1-D flat tensors.
+
+  This test case verifies that such tensors work with eager execution.
+
+  The flattening currently only happens on TPU, but tests should work
+  fine with all backends as flattening is transparent.
+  """
+
+  def testFromConstant(self):
+    with self.test_scope():
+      # Create constant of shape [100, 2, 1]. This tensor would be
+      # excessively padded on TPU.
+      tensor = constant_op.constant(100 * [[[10.0], [2.0]]])
+      # Use reduce_sum since it requires correctly working with
+      # a particular dimension.
+      reduced = math_ops.reduce_sum(tensor, axis=1)
+      self.assertAllEqual(100 * [[12.0]], reduced)
+
+  def testFromOperation(self):
+    with self.test_scope():
+      tensor = array_ops.ones([3, 100, 2, 2])
+      reduced = math_ops.reduce_sum(tensor, axis=[0, 2, 3])
+      self.assertAllEqual(100 * [12.0], reduced)
+
+  def testAsFunctionInput(self):
+    with self.test_scope():
+
+      @def_function.function
+      def f(x):
+        return math_ops.reduce_sum(x, axis=2)
+
+      tensor = constant_op.constant(100 * [[[10.0, 2.0]]])
+      reduced = f(tensor)
+      self.assertAllEqual(100 * [[12.0]], reduced)
+
+  def testAsFunctionOutput(self):
+    with self.test_scope():
+
+      @def_function.function
+      def f(x):
+        return x * constant_op.constant(100 * [[[10.0, 2.0]]])
+
+      y = f(3)
+      reduced = math_ops.reduce_sum(y, axis=2)
+      self.assertAllEqual(100 * [[36.0]], reduced)
+
+
+def multiple_tpus():
+  devices = context.context().devices()
+  return len([d for d in devices if 'device:TPU:' in d]) > 1
+
+
+class MultiDeviceTest(xla_test.XLATestCase):
+  """Test running TPU computation on more than one core."""
+
+  def testBasic(self):
+    if not multiple_tpus():
+      self.skipTest('MultiDeviceTest requires multiple TPU devices.')
+
+    # Compute 10 on TPU core 0
+    with ops.device('device:TPU:0'):
+      two = constant_op.constant(2)
+      five = constant_op.constant(5)
+      ten = two * five
+      self.assertAllEqual(10, ten)
+
+    # Compute 6 on TPU core 1
+    with ops.device('device:TPU:1'):
+      two = constant_op.constant(2)
+      three = constant_op.constant(3)
+      six = two * three
+      self.assertAllEqual(6, six)
+
+    # Copy 10 and 6 to CPU and sum them
+    self.assertAllEqual(16, ten + six)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution(
+      config=config_pb2.ConfigProto(log_device_placement=True))
+  googletest.main()

tensorflow/compiler/tests/einsum_op_test.py

@@ -0,0 +1,84 @@
+# 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.
+# ==============================================================================
+"""Test cases for einsum op."""
+
+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 special_math_ops
+from tensorflow.python.platform import googletest
+
+
+class EinsumOpTest(xla_test.XLATestCase):
+  """Test cases for einsum op."""
+
+  def _testUnary(self, op, inp, expected):
+    """Verifies that unary 'op' produces 'expected' when fed input 'inp'."""
+    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})
+      self.assertEqual(output.dtype, expected.dtype)
+      self.assertAllCloseAccordingToType(
+          expected, result, rtol=1e-3, atol=1e-5, bfloat16_rtol=0.03)
+
+  def _testBinary(self, op, a, b, expected):
+    """Verifies that binary 'op' produces 'expected' when fed 'a' and 'b'."""
+    with self.session() as session:
+      with self.test_scope():
+        pa = array_ops.placeholder(dtypes.as_dtype(a.dtype), a.shape, name='a')
+        pb = array_ops.placeholder(dtypes.as_dtype(b.dtype), b.shape, name='b')
+        output = op(pa, pb)
+      result = session.run(output, {pa: a, pb: b})
+      self.assertAllCloseAccordingToType(result, expected, rtol=1e-3)
+
+  def testMatMul(self):
+    for dtype in self.float_types:
+      self._testBinary(
+          lambda x, y: special_math_ops.einsum('ij,jk->ik', x, y),
+          np.array([[-0.25]], dtype=dtype),
+          np.array([[8]], dtype=dtype),
+          expected=np.array([[-2]], dtype=dtype))
+
+  def testImplicitForm(self):
+    for dtype in self.float_types:
+      self._testBinary(
+          lambda x, y: special_math_ops.einsum('ijk,kji', x, y),
+          np.array([[[1, 3], [2, 5], [6, 8]]], dtype=dtype),
+          np.array([[[1], [3], [2]], [[5], [6], [8]]], dtype=dtype),
+          expected=np.array(128, dtype=dtype))
+
+  def testReducedIndices(self):
+    for dtype in self.float_types:
+      self._testBinary(
+          lambda x, y: special_math_ops.einsum('ij,j->', x, y),
+          np.array([[1, 3], [2, 5], [6, 8]], dtype=dtype),
+          np.array([3, 2], dtype=dtype),
+          expected=np.array(59, dtype=dtype))
+
+  def testUnary(self):
+    for dtype in self.float_types:
+      self._testUnary(
+          lambda x: special_math_ops.einsum('ijk->kji', x),
+          np.array([[[1, 3], [2, 5], [6, 8]]], dtype=dtype),
+          expected=np.array([[[1], [2], [6]], [[3], [5], [8]]], dtype=dtype))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/ensure_shape_op_test.py

@@ -0,0 +1,47 @@
+# Copyright 2016 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 ensure_shape_op."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.platform import test
+
+
+class EnsureShapeOpTest(xla_test.XLATestCase):
+
+  def testEnsureShape(self):
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.int32)
+      with self.test_scope():
+        op = check_ops.ensure_shape(p, (None, 3))
+      expected_out = [[0, 1, 2], [3, 4, 5], [6, 7, 8]]
+      self.assertAllEqual(expected_out,
+                          sess.run(op, {p: [[0, 1, 2], [3, 4, 5], [6, 7, 8]]}))
+
+  def testInvalidEnsureShape(self):
+    with self.session() as sess:
+      p = array_ops.placeholder(dtypes.int32)
+      with self.test_scope():
+        op = check_ops.ensure_shape(p, (None, 3, 3))
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  "is not compatible with expected shape"):
+        sess.run(op, {p: [[0, 1, 2], [3, 4, 5], [6, 7, 8]]})
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/extract_image_patches_op_test.py

@@ -0,0 +1,164 @@
+# Copyright 2018 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 ExtractImagePatches op."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ExtractImagePatches(xla_test.XLATestCase):
+  """Functional tests for ExtractImagePatches op."""
+
+  def _VerifyValues(self, image, ksizes, strides, rates, padding, patches):
+    """Tests input-output pairs for the ExtractImagePatches op.
+
+    Args:
+      image: Input tensor with shape: [batch, in_rows, in_cols, depth].
+      ksizes: Patch size specified as: [ksize_rows, ksize_cols].
+      strides: Output strides, specified as [stride_rows, stride_cols].
+      rates: Atrous rates, specified as [rate_rows, rate_cols].
+      padding: Padding type.
+      patches: Expected output.
+    """
+    ksizes = [1] + ksizes + [1]
+    strides = [1] + strides + [1]
+    rates = [1] + rates + [1]
+
+    with self.session():
+      image_placeholder = array_ops.placeholder(dtypes.float32)
+      with self.test_scope():
+        out_tensor = array_ops.extract_image_patches(
+            image_placeholder,
+            ksizes=ksizes,
+            strides=strides,
+            rates=rates,
+            padding=padding,
+            name="im2col")
+      feed_dict = {image_placeholder: image}
+      self.assertAllClose(patches, out_tensor.eval(feed_dict=feed_dict))
+
+  def testKsize1x1Stride1x1Rate1x1(self):
+    """Verifies that for 1x1 kernel the output equals the input."""
+    # [2, 3, 4, 5]
+    image = np.reshape(range(120), [2, 3, 4, 5])
+    # [2, 3, 4, 5]
+    patches = np.reshape(range(120), [2, 3, 4, 5])
+    for padding in ["VALID", "SAME"]:
+      self._VerifyValues(
+          image,
+          ksizes=[1, 1],
+          strides=[1, 1],
+          rates=[1, 1],
+          padding=padding,
+          patches=patches)
+
+  def testKsize1x1Stride2x3Rate1x1(self):
+    """Test for 1x1 kernel and strides."""
+    # [2, 4, 5, 3]
+    image = np.reshape(range(120), [2, 4, 5, 3])
+    # [2, 2, 2, 3]
+    patches = image[:, ::2, ::3, :]
+    for padding in ["VALID", "SAME"]:
+      self._VerifyValues(
+          image,
+          ksizes=[1, 1],
+          strides=[2, 3],
+          rates=[1, 1],
+          padding=padding,
+          patches=patches)
+
+  def testKsize2x2Stride1x1Rate1x1Valid(self):
+    """Test for 2x2 kernel with VALID padding."""
+    # [1, 2, 2, 1]
+    image = [[[[1], [2]], [[3], [4]]]]
+    # [1, 1, 1, 4]
+    patches = [[[[1, 2, 3, 4]]]]
+    self._VerifyValues(
+        image,
+        ksizes=[2, 2],
+        strides=[1, 1],
+        rates=[1, 1],
+        padding="VALID",
+        patches=patches)
+
+  def testKsize2x2Stride1x1Rate1x1Same(self):
+    """Test for 2x2 kernel with SAME padding."""
+    # [1, 2, 2, 1]
+    image = [[[[1], [2]], [[3], [4]]]]
+    # [1, 2, 2, 4]
+    patches = [[[[1, 2, 3, 4], [2, 0, 4, 0]], [[3, 4, 0, 0], [4, 0, 0, 0]]]]
+    self._VerifyValues(
+        image,
+        ksizes=[2, 2],
+        strides=[1, 1],
+        rates=[1, 1],
+        padding="SAME",
+        patches=patches)
+
+  def testKsize2x2Stride1x1Rate2x2Valid(self):
+    """Test for 2x2 kernel with 2x2 dilation."""
+    # [1, 2, 2, 1]
+    image = np.arange(16).reshape(1, 4, 4, 1).astype(np.float32)
+    # [1, 2, 2, 4]
+    patches = [[[[0, 2, 8, 10], [1, 3, 9, 11]],
+                [[4, 6, 12, 14], [5, 7, 13, 15]]]]
+    self._VerifyValues(
+        image,
+        ksizes=[2, 2],
+        strides=[1, 1],
+        rates=[2, 2],
+        padding="VALID",
+        patches=patches)
+
+  def testKsize2x2Stride1x1Rate1x1ValidDepth2(self):
+    """Test for 2x2 kernel with VALID padding."""
+    # [1, 2, 2, 2]
+    image = [[[[1, 5], [2, 6]], [[3, 7], [4, 8]]]]
+    # [1, 1, 1, 8]
+    patches = [[[[1, 5, 2, 6, 3, 7, 4, 8]]]]
+    self._VerifyValues(
+        image,
+        ksizes=[2, 2],
+        strides=[1, 1],
+        rates=[1, 1],
+        padding="VALID",
+        patches=patches)
+
+  def testInvalidKernelSize(self):
+    """Test that zero kernel size raises an error in XLA mode."""
+    with self.session():
+      image_placeholder = array_ops.placeholder(dtypes.float32)
+      with self.test_scope():
+        out_tensor = array_ops.extract_image_patches(
+            image_placeholder,
+            ksizes=[1, 0, 2, 1],
+            strides=[1, 1, 1, 1],
+            rates=[1, 1, 1, 1],
+            padding="VALID",
+        )
+      feed_dict = {image_placeholder: np.zeros([1, 4, 4, 1])}
+      with self.assertRaisesRegex(
+          errors_impl.OutOfRangeError, "Kernel size values must be positive"
+      ):
+        out_tensor.eval(feed_dict=feed_dict)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/fake_quant_ops_test.py

@@ -0,0 +1,695 @@
+# Copyright 2018 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.
+# ==============================================================================
+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 gen_array_ops
+from tensorflow.python.platform import googletest
+
+
+class FakeQuantWithMinMaxArgsTest(xla_test.XLATestCase):
+  """Test cases for FakeQuantWithMinMaxArgs operation."""
+
+  # 8 bits, wide range.
+  def testOp_with8BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0)
+
+  def testOp_with8BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingUp(self):
+    self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5)
+
+  def testOp_with8BitsRoundingHalfWayZeroPoint(self):
+    # zero_point_from_min lands at exactly 0.5 (half-integer boundary).
+    # Scale = (509.0 - (-1.0)) / (255 - 0) = 510.0 / 255.0 = 2.0
+    # zero_point_from_min = 0 - (-1.0) / 2.0 = 0.5
+    # Under round-half-away-from-zero: nudged_zero_point = 1.0
+    # nudged_min = (0 - 1.0) * 2.0 = -2.0
+    # nudged_max = (255 - 1.0) * 2.0 = 508.0
+    self._TestOp(-1.0, 509.0, 8, False, -2.0, 508.0, 2.0)
+
+  # 8 bits, narrow range.
+  def testOp_with8BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5)
+
+  # 7 bits, wide range.
+  def testOp_with7BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0)
+
+  def testOp_with7BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingUp(self):
+    self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5)
+
+  # 7 bits, narrow range.
+  def testOp_with7BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5)
+
+  def _TestOp(self, input_min, input_max, num_bits, narrow_range,
+              expected_nudged_input_min, expected_nudged_input_max,
+              expected_step):
+    inputs = np.array(
+        [
+            expected_nudged_input_min - expected_step,
+            expected_nudged_input_min - 0.01, expected_nudged_input_min,
+            expected_nudged_input_min + 0.01,
+            expected_nudged_input_min + expected_step - 0.01,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step + 0.01,
+            expected_nudged_input_max - 0.01, expected_nudged_input_max,
+            expected_nudged_input_max + 0.01,
+            expected_nudged_input_max + expected_step
+        ],
+        dtype=np.float32)
+    expected = np.array(
+        [
+            expected_nudged_input_min, expected_nudged_input_min,
+            expected_nudged_input_min, expected_nudged_input_min,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_max, expected_nudged_input_max,
+            expected_nudged_input_max, expected_nudged_input_max
+        ],
+        dtype=np.float32)
+
+    with self.session() as session:
+      with self.test_scope():
+        input_placeholder = array_ops.placeholder(
+            dtypes.float32, inputs.shape, name="inputs")
+        outputs = array_ops.fake_quant_with_min_max_args(
+            input_placeholder,
+            min=input_min,
+            max=input_max,
+            num_bits=num_bits,
+            narrow_range=narrow_range)
+      result = session.run(outputs, {input_placeholder: inputs})
+      self.assertAllCloseAccordingToType(
+          result, expected, rtol=1e-3, atol=1e-5, bfloat16_rtol=0.03)
+
+
+class FakeQuantWithMinMaxArgsGradientTest(xla_test.XLATestCase):
+  """Test cases for FakeQuantWithMinMaxArgsGradient operation."""
+
+  # 8 bits, wide range.
+  def testOp_with8BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0)
+
+  def testOp_with8BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingUp(self):
+    self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5)
+
+  # 8 bits, narrow range.
+  def testOp_with8BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5)
+
+  # 7 bits, wide range.
+  def testOp_with7BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0)
+
+  def testOp_with7BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingUp(self):
+    self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5)
+
+  # 7 bits, narrow range.
+  def testOp_with7BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5)
+
+  def _TestOp(self, input_min, input_max, num_bits, narrow_range,
+              expected_nudged_input_min, expected_nudged_input_max,
+              expected_step):
+    inputs = np.array(
+        [
+            expected_nudged_input_min - expected_step,
+            expected_nudged_input_min - 0.01, expected_nudged_input_min,
+            expected_nudged_input_min + 0.01,
+            expected_nudged_input_min + expected_step - 0.01,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step + 0.01,
+            expected_nudged_input_max - 0.01, expected_nudged_input_max,
+            expected_nudged_input_max + 0.01,
+            expected_nudged_input_max + expected_step
+        ],
+        dtype=np.float32)
+    gradients = np.arange(1, len(inputs) + 1, dtype=np.float32)
+    expected_backprops = np.array(
+        [0.0, 0.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 0.0, 0.0],
+        dtype=np.float32)
+
+    with self.session() as session:
+      with self.test_scope():
+        gradient_placeholder = array_ops.placeholder(
+            dtypes.float32, gradients.shape, name="gradients")
+        input_placeholder = array_ops.placeholder(
+            dtypes.float32, inputs.shape, name="inputs")
+        outputs = gen_array_ops.fake_quant_with_min_max_args_gradient(
+            gradient_placeholder,
+            input_placeholder,
+            min=input_min,
+            max=input_max,
+            num_bits=num_bits,
+            narrow_range=narrow_range)
+      backprops = session.run(outputs, {
+          gradient_placeholder: gradients,
+          input_placeholder: inputs
+      })
+      self.assertAllCloseAccordingToType(
+          backprops,
+          expected_backprops,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+
+
+class FakeQuantWithMinMaxVarsTest(xla_test.XLATestCase):
+  """Test cases for FakeQuantWithMinMaxVars operation."""
+
+  # 8 bits, wide range.
+  def testOp_with8BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0)
+
+  def testOp_with8BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingUp(self):
+    self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5)
+
+  # 8 bits, narrow range.
+  def testOp_with8BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5)
+
+  # 7 bits, wide range.
+  def testOp_with7BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0)
+
+  def testOp_with7BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingUp(self):
+    self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5)
+
+  # 7 bits, narrow range.
+  def testOp_with7BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5)
+
+  def _TestOp(self, input_min, input_max, num_bits, narrow_range,
+              expected_nudged_input_min, expected_nudged_input_max,
+              expected_step):
+    inputs = np.array(
+        [
+            expected_nudged_input_min - expected_step,
+            expected_nudged_input_min - 0.01, expected_nudged_input_min,
+            expected_nudged_input_min + 0.01,
+            expected_nudged_input_min + expected_step - 0.01,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step + 0.01,
+            expected_nudged_input_max - 0.01, expected_nudged_input_max,
+            expected_nudged_input_max + 0.01,
+            expected_nudged_input_max + expected_step
+        ],
+        dtype=np.float32)
+    expected = np.array(
+        [
+            expected_nudged_input_min, expected_nudged_input_min,
+            expected_nudged_input_min, expected_nudged_input_min,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_max, expected_nudged_input_max,
+            expected_nudged_input_max, expected_nudged_input_max
+        ],
+        dtype=np.float32)
+
+    with self.session() as session:
+      with self.test_scope():
+        input_placeholder = array_ops.placeholder(
+            dtypes.float32, inputs.shape, name="inputs")
+        min_placeholder = array_ops.placeholder(dtypes.float32, (), name="min")
+        max_placeholder = array_ops.placeholder(dtypes.float32, (), name="max")
+        outputs = array_ops.fake_quant_with_min_max_vars(
+            input_placeholder,
+            min_placeholder,
+            max_placeholder,
+            num_bits=num_bits,
+            narrow_range=narrow_range)
+      result = session.run(
+          outputs, {
+              input_placeholder: inputs,
+              min_placeholder: input_min,
+              max_placeholder: input_max
+          })
+      self.assertAllCloseAccordingToType(
+          result, expected, rtol=1e-3, atol=1e-5, bfloat16_rtol=0.03)
+
+
+class FakeQuantWithMinMaxVarsGradientTest(xla_test.XLATestCase):
+  """Test cases for FakeQuantWithMinMaxVarsGradient operation."""
+
+  # 8 bits, wide range.
+  def testOp_with8BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 255.0, 8, False, 0.0, 255.0, 1.0)
+
+  def testOp_with8BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 128.0, 8, False, 0.0, 127.5, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingUp(self):
+    self._TestOp(-128.0, -0.5, 8, False, -127.5, 0.0, 0.5)
+
+  def testOp_with8BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 127.4, 8, False, 0.0, 127.5, 0.5)
+
+  # 8 bits, narrow range.
+  def testOp_with8BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 254.0, 8, True, 0.0, 254.0, 1.0)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 127.1, 8, True, 0.0, 127.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-127.1, -0.1, 8, True, -127.0, 0.0, 0.5)
+
+  def testOp_with8BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 126.9, 8, True, 0.0, 127.0, 0.5)
+
+  # 7 bits, wide range.
+  def testOp_with7BitsNoScalingNoNudging(self):
+    self._TestOp(0.0, 127.0, 7, False, 0.0, 127.0, 1.0)
+
+  def testOp_with7BitsScalingAndNudgingDown(self):
+    self._TestOp(0.5, 64.0, 7, False, 0.0, 63.5, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingUp(self):
+    self._TestOp(-64.0, -0.5, 7, False, -63.5, 0.0, 0.5)
+
+  def testOp_with7BitsScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 63.4, 7, False, 0.0, 63.5, 0.5)
+
+  # 7 bits, narrow range.
+  def testOp_with7BitsNarrowRangeNoScalingNoNudging(self):
+    self._TestOp(0.0, 126.0, 7, True, 0.0, 126.0, 1.0)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingDown(self):
+    self._TestOp(0.1, 63.1, 7, True, 0.0, 63.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingUp(self):
+    self._TestOp(-63.1, -0.1, 7, True, -63.0, 0.0, 0.5)
+
+  def testOp_with7BitsNarrowRangeScalingAndNudgingBetween(self):
+    self._TestOp(-0.1, 62.9, 7, True, 0.0, 63.0, 0.5)
+
+  def _TestOp(self, input_min, input_max, num_bits, narrow_range,
+              expected_nudged_input_min, expected_nudged_input_max,
+              expected_step):
+    inputs = np.array(
+        [
+            expected_nudged_input_min - expected_step,
+            expected_nudged_input_min - 0.01, expected_nudged_input_min,
+            expected_nudged_input_min + 0.01,
+            expected_nudged_input_min + expected_step - 0.01,
+            expected_nudged_input_min + expected_step,
+            expected_nudged_input_min + expected_step + 0.01,
+            expected_nudged_input_max - 0.01, expected_nudged_input_max,
+            expected_nudged_input_max + 0.01,
+            expected_nudged_input_max + expected_step
+        ],
+        dtype=np.float32)
+    gradients = np.arange(1, len(inputs) + 1, dtype=np.float32)
+    expected_backprops_wrt_input = np.array(
+        [0.0, 0.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 0.0, 0.0],
+        dtype=np.float32)
+    expected_backprops_wrt_min = 1.0 + 2.0
+    expected_backprops_wrt_max = 10.0 + 11.0
+
+    with self.session() as session:
+      with self.test_scope():
+        gradient_placeholder = array_ops.placeholder(
+            dtypes.float32, gradients.shape, name="gradients")
+        input_placeholder = array_ops.placeholder(
+            dtypes.float32, inputs.shape, name="inputs")
+        min_placeholder = array_ops.placeholder(dtypes.float32, (), name="min")
+        max_placeholder = array_ops.placeholder(dtypes.float32, (), name="max")
+        outputs = array_ops.fake_quant_with_min_max_vars_gradient(
+            gradient_placeholder,
+            input_placeholder,
+            min_placeholder,
+            max_placeholder,
+            num_bits=num_bits,
+            narrow_range=narrow_range)
+      backprops_wrt_input, backprops_wrt_min, backprops_wrt_max = session.run(
+          outputs, {
+              gradient_placeholder: gradients,
+              input_placeholder: inputs,
+              min_placeholder: input_min,
+              max_placeholder: input_max
+          })
+      self.assertAllCloseAccordingToType(
+          backprops_wrt_input,
+          expected_backprops_wrt_input,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+      self.assertAllCloseAccordingToType(
+          backprops_wrt_min,
+          expected_backprops_wrt_min,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+      self.assertAllCloseAccordingToType(
+          backprops_wrt_max,
+          expected_backprops_wrt_max,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+
+
+class FakeQuantWithMinMaxVarsPerChannelTest(xla_test.XLATestCase):
+  """Test cases for FakeQuantWithMinMaxVarsPerChannel operation."""
+
+  # 8 bits, wide range.
+  def testOp_with8Bits(self):
+    self._TestOp(
+        [0.0, 0.5, -128.0, -0.1],
+        [255.0, 128.0, -0.5, 127.4],
+        8,
+        False,
+        [0.0, 0.0, -127.5, 0.0],
+        [255.0, 127.5, 0.0, 127.5],
+        [1.0, 0.5, 0.5, 0.5])
+
+  # 8 bits, narrow range.
+  def testOp_with8BitsNarrowRange(self):
+    self._TestOp(
+        [0.0, 0.1, -127.1, -0.1],
+        [254.0, 127.1, -0.1, 126.9],
+        8,
+        True,
+        [0.0, 0.0, -127.0, 0.0],
+        [254.0, 127.0, 0.0, 127.0],
+        [1.0, 0.5, 0.5, 0.5])
+
+  # 7 bits, wide range.
+  def testOp_with7Bits(self):
+    self._TestOp(
+        [0.0, 0.5, -64.0, -0.1],
+        [127.0, 64.0, -0.5, 63.4],
+        7,
+        False,
+        [0.0, 0.0, -63.5, 0.0],
+        [127.0, 63.5, 0.0, 63.5],
+        [1.0, 0.5, 0.5, 0.5])
+
+  # 7 bits, narrow range.
+  def testOp_with7BitsNarrowRange(self):
+    self._TestOp(
+        [0.0, 0.1, -63.1, -0.1],
+        [126.0, 63.1, -0.1, 62.9],
+        7,
+        True,
+        [0.0, 0.0, -63.0, 0.0],
+        [126.0, 63.0, 0.0, 63.0],
+        [1.0, 0.5, 0.5, 0.5])
+
+  def _TestOp(self, input_mins, input_maxs, num_bits, narrow_range,
+              expected_nudged_input_mins, expected_nudged_input_maxs,
+              expected_steps):
+    num_channels = len(input_mins)
+    inputs_list = []
+    expected_list = []
+    for i in range(num_channels):
+      expected_nudged_input_min = expected_nudged_input_mins[i]
+      expected_nudged_input_max = expected_nudged_input_maxs[i]
+      expected_step = expected_steps[i]
+
+      inputs_list.append(
+          [
+              expected_nudged_input_min - expected_step,
+              expected_nudged_input_min - 0.01, expected_nudged_input_min,
+              expected_nudged_input_min + 0.01,
+              expected_nudged_input_min + expected_step - 0.01,
+              expected_nudged_input_min + expected_step,
+              expected_nudged_input_min + expected_step + 0.01,
+              expected_nudged_input_max - 0.01, expected_nudged_input_max,
+              expected_nudged_input_max + 0.01,
+              expected_nudged_input_max + expected_step
+          ])
+      expected_list.append(
+          [
+              expected_nudged_input_min, expected_nudged_input_min,
+              expected_nudged_input_min, expected_nudged_input_min,
+              expected_nudged_input_min + expected_step,
+              expected_nudged_input_min + expected_step,
+              expected_nudged_input_min + expected_step,
+              expected_nudged_input_max, expected_nudged_input_max,
+              expected_nudged_input_max, expected_nudged_input_max
+          ])
+    inputs = np.transpose(np.array(inputs_list, dtype=np.float32))
+    expected = np.transpose(np.array(expected_list, dtype=np.float32))
+
+    with self.session() as session:
+      with self.test_scope():
+        input_placeholder = array_ops.placeholder(
+            dtypes.float32, inputs.shape, name="inputs")
+        min_placeholder = array_ops.placeholder(
+            dtypes.float32, (num_channels), name="min")
+        max_placeholder = array_ops.placeholder(
+            dtypes.float32, (num_channels), name="max")
+        outputs = array_ops.fake_quant_with_min_max_vars_per_channel(
+            input_placeholder,
+            min_placeholder,
+            max_placeholder,
+            num_bits=num_bits,
+            narrow_range=narrow_range)
+      result = session.run(
+          outputs, {
+              input_placeholder: inputs,
+              min_placeholder: input_mins,
+              max_placeholder: input_maxs
+          })
+      self.assertAllCloseAccordingToType(
+          result, expected, rtol=1e-3, atol=1e-5, bfloat16_rtol=0.03)
+
+
+class FakeQuantWithMinMaxVarsPerChannelGradientTest(xla_test.XLATestCase):
+  """Test cases for FakeQuantWithMinMaxVarsPerChannelGradient operation."""
+
+  # 8 bits, wide range.
+  def testOp_with8Bits(self):
+    self._TestOp(
+        [0.0, 0.5, -128.0, -0.1],
+        [255.0, 128.0, -0.5, 127.4],
+        8,
+        False,
+        [0.0, 0.0, -127.5, 0.0],
+        [255.0, 127.5, 0.0, 127.5],
+        [1.0, 0.5, 0.5, 0.5])
+
+  # 8 bits, narrow range.
+  def testOp_with8BitsNarrowRange(self):
+    self._TestOp(
+        [0.0, 0.1, -127.1, -0.1],
+        [254.0, 127.1, -0.1, 126.9],
+        8,
+        True,
+        [0.0, 0.0, -127.0, 0.0],
+        [254.0, 127.0, 0.0, 127.0],
+        [1.0, 0.5, 0.5, 0.5])
+
+  # 7 bits, wide range.
+  def testOp_with7Bits(self):
+    self._TestOp(
+        [0.0, 0.5, -64.0, -0.1],
+        [127.0, 64.0, -0.5, 63.4],
+        7,
+        False,
+        [0.0, 0.0, -63.5, 0.0],
+        [127.0, 63.5, 0.0, 63.5],
+        [1.0, 0.5, 0.5, 0.5])
+
+  # 7 bits, narrow range.
+  def testOp_with7BitsNarrowRange(self):
+    self._TestOp(
+        [0.0, 0.1, -63.1, -0.1],
+        [126.0, 63.1, -0.1, 62.9],
+        7,
+        True,
+        [0.0, 0.0, -63.0, 0.0],
+        [126.0, 63.0, 0.0, 63.0],
+        [1.0, 0.5, 0.5, 0.5])
+
+  def _TestOp(self, input_mins, input_maxs, num_bits, narrow_range,
+              expected_nudged_input_mins, expected_nudged_input_maxs,
+              expected_steps):
+    num_channels = len(input_mins)
+    inputs_list = []
+    gradients_list = []
+    expected_backprops_wrt_input_list = []
+    expected_backprops_wrt_min_list = []
+    expected_backprops_wrt_max_list = []
+    for i in range(num_channels):
+      expected_nudged_input_min = expected_nudged_input_mins[i]
+      expected_nudged_input_max = expected_nudged_input_maxs[i]
+      expected_step = expected_steps[i]
+      inputs = [
+          expected_nudged_input_min - expected_step,
+          expected_nudged_input_min - 0.01, expected_nudged_input_min,
+          expected_nudged_input_min + 0.01,
+          expected_nudged_input_min + expected_step - 0.01,
+          expected_nudged_input_min + expected_step,
+          expected_nudged_input_min + expected_step + 0.01,
+          expected_nudged_input_max - 0.01, expected_nudged_input_max,
+          expected_nudged_input_max + 0.01,
+          expected_nudged_input_max + expected_step
+      ]
+      inputs_list.append(inputs)
+      gradients_list.append(list(range(1, len(inputs) + 1)))
+      expected_backprops_wrt_input_list.append(
+          [0.0, 0.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 0.0, 0.0])
+      expected_backprops_wrt_min_list.append(1.0 + 2.0)
+      expected_backprops_wrt_max_list.append(10.0 + 11.0)
+
+    inputs = np.transpose(np.array(inputs_list, dtype=np.float32))
+    gradients = np.transpose(np.array(gradients_list, dtype=np.float32))
+    expected_backprops_wrt_input = np.transpose(np.array(
+        expected_backprops_wrt_input_list, dtype=np.float32))
+    expected_backprops_wrt_min = np.array(
+        expected_backprops_wrt_min_list, dtype=np.float32)
+    expected_backprops_wrt_max = np.array(
+        expected_backprops_wrt_max_list, dtype=np.float32)
+
+    with self.session() as session:
+      with self.test_scope():
+        gradient_placeholder = array_ops.placeholder(
+            dtypes.float32, gradients.shape, name="gradients")
+        input_placeholder = array_ops.placeholder(
+            dtypes.float32, inputs.shape, name="inputs")
+        min_placeholder = array_ops.placeholder(
+            dtypes.float32, (num_channels), name="min")
+        max_placeholder = array_ops.placeholder(
+            dtypes.float32, (num_channels), name="max")
+        outputs = array_ops.fake_quant_with_min_max_vars_per_channel_gradient(
+            gradient_placeholder,
+            input_placeholder,
+            min_placeholder,
+            max_placeholder,
+            num_bits=num_bits,
+            narrow_range=narrow_range)
+      backprops_wrt_input, backprops_wrt_min, backprops_wrt_max = session.run(
+          outputs, {
+              gradient_placeholder: gradients,
+              input_placeholder: inputs,
+              min_placeholder: input_mins,
+              max_placeholder: input_maxs
+          })
+      self.assertAllCloseAccordingToType(
+          backprops_wrt_input,
+          expected_backprops_wrt_input,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+      self.assertAllCloseAccordingToType(
+          backprops_wrt_min,
+          expected_backprops_wrt_min,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+      self.assertAllCloseAccordingToType(
+          backprops_wrt_max,
+          expected_backprops_wrt_max,
+          rtol=1e-3,
+          atol=1e-5,
+          bfloat16_rtol=0.03)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/fft_test.py

@@ -0,0 +1,247 @@
+# 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.
+# ==============================================================================
+"""Tests for FFT via the XLA JIT."""
+
+import itertools
+
+import numpy as np
+import scipy.signal as sps
+
+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 gradients_impl
+from tensorflow.python.ops.signal import signal
+from tensorflow.python.platform import googletest
+
+BATCH_DIMS = (3, 5)
+RTOL = 0.009  # Eigen/cuFFT differ widely from np, especially for FFT3D
+ATOL = 1e-4
+RTOL_3D = 0.07
+ATOL_3D = 4e-4
+
+
+def pick_10(x):
+  x = list(x)
+  np.random.seed(123)
+  np.random.shuffle(x)
+  return x[:10]
+
+
+def to_32bit(x):
+  if x.dtype == np.complex128:
+    return x.astype(np.complex64)
+  if x.dtype == np.float64:
+    return x.astype(np.float32)
+  return x
+
+
+POWS_OF_2 = 2**np.arange(3, 12)
+INNER_DIMS_1D = list((x,) for x in POWS_OF_2)
+POWS_OF_2 = 2**np.arange(3, 8)  # To avoid OOM on GPU.
+INNER_DIMS_2D = pick_10(itertools.product(POWS_OF_2, POWS_OF_2))
+INNER_DIMS_3D = pick_10(itertools.product(POWS_OF_2, POWS_OF_2, POWS_OF_2))
+
+
+class FFTTest(xla_test.XLATestCase):
+
+  def _VerifyFftMethod(self,
+                       inner_dims,
+                       complex_to_input,
+                       input_to_expected,
+                       tf_method,
+                       atol=ATOL,
+                       rtol=RTOL):
+    for indims in inner_dims:
+      print("nfft =", indims)
+      shape = BATCH_DIMS + indims
+      data = np.arange(np.prod(shape) * 2) / np.prod(indims)
+      np.random.seed(123)
+      np.random.shuffle(data)
+      data = np.reshape(data.astype(np.float32).view(np.complex64), shape)
+      data = to_32bit(complex_to_input(data))
+      expected = to_32bit(input_to_expected(data))
+      with self.session() as sess:
+        with self.test_scope():
+          ph = array_ops.placeholder(
+              dtypes.as_dtype(data.dtype), shape=data.shape)
+          out = tf_method(ph)
+        value = sess.run(out, {ph: data})
+        self.assertAllClose(expected, value, rtol=rtol, atol=atol)
+
+  def testContribSignalSTFT(self):
+    ws = 512
+    hs = 128
+    dims = (ws * 20,)
+    shape = BATCH_DIMS + dims
+    data = np.arange(np.prod(shape)) / np.prod(dims)
+    np.random.seed(123)
+    np.random.shuffle(data)
+    data = np.reshape(data.astype(np.float32), shape)
+    window = sps.get_window("hann", ws)
+    expected = sps.stft(
+        data, nperseg=ws, noverlap=ws - hs, boundary=None, window=window)[2]
+    expected = np.swapaxes(expected, -1, -2)
+    expected *= window.sum()  # scipy divides by window sum
+    with self.session() as sess:
+      with self.test_scope():
+        ph = array_ops.placeholder(
+            dtypes.as_dtype(data.dtype), shape=data.shape)
+        out = signal.stft(ph, ws, hs)
+        grad = gradients_impl.gradients(out, ph,
+                                        grad_ys=array_ops.ones_like(out))
+
+      # For gradients, we simply verify that they compile & execute.
+      value, _ = sess.run([out, grad], {ph: data})
+      self.assertAllClose(expected, value, rtol=RTOL, atol=ATOL)
+
+  def testFFT(self):
+    self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.fft,
+                          signal.fft)
+
+  def testFFT2D(self):
+    self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.fft2,
+                          signal.fft2d)
+
+  def testFFT3D(self):
+    self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x,
+                          lambda x: np.fft.fftn(x, axes=(-3, -2, -1)),
+                          signal.fft3d, ATOL_3D, RTOL_3D)
+
+  def testIFFT(self):
+    self._VerifyFftMethod(INNER_DIMS_1D, lambda x: x, np.fft.ifft,
+                          signal.ifft)
+
+  def testIFFT2D(self):
+    self._VerifyFftMethod(INNER_DIMS_2D, lambda x: x, np.fft.ifft2,
+                          signal.ifft2d)
+
+  def testIFFT3D(self):
+    self._VerifyFftMethod(INNER_DIMS_3D, lambda x: x,
+                          lambda x: np.fft.ifftn(x, axes=(-3, -2, -1)),
+                          signal.ifft3d, ATOL_3D, RTOL_3D)
+
+  def testRFFT(self):
+
+    def _to_expected(x):
+      return np.fft.rfft(x, n=x.shape[-1])
+
+    def _tf_fn(x):
+      return signal.rfft(x, fft_length=[x.shape[-1]])
+
+    self._VerifyFftMethod(INNER_DIMS_1D, np.real, _to_expected, _tf_fn)
+
+  def testRFFT2D(self):
+
+    def _tf_fn(x):
+      return signal.rfft2d(x, fft_length=[x.shape[-2], x.shape[-1]])
+
+    self._VerifyFftMethod(
+        INNER_DIMS_2D, np.real,
+        lambda x: np.fft.rfft2(x, s=[x.shape[-2], x.shape[-1]]), _tf_fn)
+
+  def testRFFT3D(self):
+
+    def _to_expected(x):
+      return np.fft.rfftn(
+          x, axes=(-3, -2, -1), s=[x.shape[-3], x.shape[-2], x.shape[-1]])
+
+    def _tf_fn(x):
+      return signal.rfft3d(
+          x, fft_length=[x.shape[-3], x.shape[-2], x.shape[-1]])
+
+    self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn, ATOL_3D,
+                          RTOL_3D)
+
+  def testRFFT3DMismatchedSize(self):
+
+    def _to_expected(x):
+      return np.fft.rfftn(
+          x,
+          axes=(-3, -2, -1),
+          s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2])
+
+    def _tf_fn(x):
+      return signal.rfft3d(
+          x, fft_length=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2])
+
+    self._VerifyFftMethod(INNER_DIMS_3D, np.real, _to_expected, _tf_fn)
+
+  def testIRFFT(self):
+
+    def _tf_fn(x):
+      return signal.irfft(x, fft_length=[2 * (x.shape[-1] - 1)])
+
+    self._VerifyFftMethod(
+        INNER_DIMS_1D, lambda x: np.fft.rfft(np.real(x), n=x.shape[-1]),
+        lambda x: np.fft.irfft(x, n=2 * (x.shape[-1] - 1)), _tf_fn)
+
+  def testIRFFT2D(self):
+
+    def _tf_fn(x):
+      return signal.irfft2d(x, fft_length=[x.shape[-2], 2 * (x.shape[-1] - 1)])
+
+    self._VerifyFftMethod(
+        INNER_DIMS_2D,
+        lambda x: np.fft.rfft2(np.real(x), s=[x.shape[-2], x.shape[-1]]),
+        lambda x: np.fft.irfft2(x, s=[x.shape[-2], 2 * (x.shape[-1] - 1)]),
+        _tf_fn)
+
+  def testIRFFT3D(self):
+
+    def _to_input(x):
+      return np.fft.rfftn(
+          np.real(x),
+          axes=(-3, -2, -1),
+          s=[x.shape[-3], x.shape[-2], x.shape[-1]])
+
+    def _to_expected(x):
+      return np.fft.irfftn(
+          x,
+          axes=(-3, -2, -1),
+          s=[x.shape[-3], x.shape[-2], 2 * (x.shape[-1] - 1)])
+
+    def _tf_fn(x):
+      return signal.irfft3d(
+          x, fft_length=[x.shape[-3], x.shape[-2], 2 * (x.shape[-1] - 1)])
+
+    self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn,
+                          ATOL_3D, RTOL_3D)
+
+  def testIRFFT3DMismatchedSize(self):
+
+    def _to_input(x):
+      return np.fft.rfftn(
+          np.real(x),
+          axes=(-3, -2, -1),
+          s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2])
+
+    def _to_expected(x):
+      return np.fft.irfftn(
+          x,
+          axes=(-3, -2, -1),
+          s=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2])
+
+    def _tf_fn(x):
+      return signal.irfft3d(
+          x, fft_length=[x.shape[-3] // 2, x.shape[-2], x.shape[-1] * 2])
+
+    self._VerifyFftMethod(INNER_DIMS_3D, _to_input, _to_expected, _tf_fn,
+                          ATOL_3D, RTOL_3D)
+
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/fifo_queue_test.py

@@ -0,0 +1,195 @@
+# Copyright 2015 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 tensorflow.ops.data_flow_ops.FIFOQueue."""
+
+import time
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes as dtypes_lib
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.platform import test
+
+
+class FIFOQueueTest(xla_test.XLATestCase):
+
+  def testEnqueue(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      enqueue_op = q.enqueue((10.0,))
+      enqueue_op.run()
+
+  def testEnqueueWithShape(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32, shapes=(3, 2))
+      enqueue_correct_op = q.enqueue(([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]],))
+      enqueue_correct_op.run()
+      with self.assertRaises(ValueError):
+        q.enqueue(([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],))
+      self.assertEqual(1, self.evaluate(q.size()))
+
+  def testMultipleDequeues(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, [dtypes_lib.int32], shapes=[()])
+      self.evaluate(q.enqueue([1]))
+      self.evaluate(q.enqueue([2]))
+      self.evaluate(q.enqueue([3]))
+      a, b, c = self.evaluate([q.dequeue(), q.dequeue(), q.dequeue()])
+      self.assertAllEqual(set([1, 2, 3]), set([a, b, c]))
+
+  def testQueuesDontShare(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, [dtypes_lib.int32], shapes=[()])
+      self.evaluate(q.enqueue(1))
+      q2 = data_flow_ops.FIFOQueue(10, [dtypes_lib.int32], shapes=[()])
+      self.evaluate(q2.enqueue(2))
+      self.assertAllEqual(self.evaluate(q2.dequeue()), 2)
+      self.assertAllEqual(self.evaluate(q.dequeue()), 1)
+
+  def testEnqueueDictWithoutNames(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      with self.assertRaisesRegex(ValueError, "must have names"):
+        q.enqueue({"a": 12.0})
+
+  def testParallelEnqueue(self):
+    with self.session() as sess, self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      elems = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]
+      enqueue_ops = [q.enqueue((x,)) for x in elems]
+      dequeued_t = q.dequeue()
+
+      # Run one producer thread for each element in elems.
+      def enqueue(enqueue_op):
+        sess.run(enqueue_op)
+
+      threads = [
+          self.checkedThread(target=enqueue, args=(e,)) for e in enqueue_ops
+      ]
+      for thread in threads:
+        thread.start()
+      for thread in threads:
+        thread.join()
+
+      # Dequeue every element using a single thread.
+      results = []
+      for _ in range(len(elems)):
+        results.append(self.evaluate(dequeued_t))
+      self.assertItemsEqual(elems, results)
+
+  def testParallelDequeue(self):
+    with self.session() as sess, self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      elems = [10.0, 20.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0, 90.0, 100.0]
+      enqueue_ops = [q.enqueue((x,)) for x in elems]
+      dequeued_t = q.dequeue()
+
+      # Enqueue every element using a single thread.
+      for enqueue_op in enqueue_ops:
+        enqueue_op.run()
+
+      # Run one consumer thread for each element in elems.
+      results = []
+
+      def dequeue():
+        results.append(sess.run(dequeued_t))
+
+      threads = [self.checkedThread(target=dequeue) for _ in enqueue_ops]
+      for thread in threads:
+        thread.start()
+      for thread in threads:
+        thread.join()
+      self.assertItemsEqual(elems, results)
+
+  def testDequeue(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      elems = [10.0, 20.0, 30.0]
+      enqueue_ops = [q.enqueue((x,)) for x in elems]
+      dequeued_t = q.dequeue()
+
+      for enqueue_op in enqueue_ops:
+        enqueue_op.run()
+
+      for i in range(len(elems)):
+        vals = self.evaluate(dequeued_t)
+        self.assertEqual([elems[i]], vals)
+
+  def testEnqueueAndBlockingDequeue(self):
+    with self.session() as sess, self.test_scope():
+      q = data_flow_ops.FIFOQueue(3, dtypes_lib.float32)
+      elems = [10.0, 20.0, 30.0]
+      enqueue_ops = [q.enqueue((x,)) for x in elems]
+      dequeued_t = q.dequeue()
+
+      def enqueue():
+        # The enqueue_ops should run after the dequeue op has blocked.
+        # TODO(mrry): Figure out how to do this without sleeping.
+        time.sleep(0.1)
+        for enqueue_op in enqueue_ops:
+          sess.run(enqueue_op)
+
+      results = []
+
+      def dequeue():
+        for _ in range(len(elems)):
+          results.append(sess.run(dequeued_t))
+
+      enqueue_thread = self.checkedThread(target=enqueue)
+      dequeue_thread = self.checkedThread(target=dequeue)
+      enqueue_thread.start()
+      dequeue_thread.start()
+      enqueue_thread.join()
+      dequeue_thread.join()
+
+      for elem, result in zip(elems, results):
+        self.assertEqual([elem], result)
+
+  def testMultiEnqueueAndDequeue(self):
+    with self.session() as sess, self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, (dtypes_lib.int32, dtypes_lib.float32))
+      elems = [(5, 10.0), (10, 20.0), (15, 30.0)]
+      enqueue_ops = [q.enqueue((x, y)) for x, y in elems]
+      dequeued_t = q.dequeue()
+
+      for enqueue_op in enqueue_ops:
+        enqueue_op.run()
+
+      for i in range(len(elems)):
+        x_val, y_val = sess.run(dequeued_t)
+        x, y = elems[i]
+        self.assertEqual([x], x_val)
+        self.assertEqual([y], y_val)
+
+  def testQueueSizeEmpty(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      self.assertEqual([0], self.evaluate(q.size()))
+
+  def testQueueSizeAfterEnqueueAndDequeue(self):
+    with self.session(), self.test_scope():
+      q = data_flow_ops.FIFOQueue(10, dtypes_lib.float32)
+      enqueue_op = q.enqueue((10.0,))
+      dequeued_t = q.dequeue()
+      size = q.size()
+      self.assertEqual([], size.get_shape())
+
+      enqueue_op.run()
+      self.assertEqual(1, self.evaluate(size))
+      dequeued_t.op.run()
+      self.assertEqual(0, self.evaluate(size))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/float_ops_test.py

@@ -0,0 +1,545 @@
+# Copyright 2025 The OpenXLA Authors.
+#
+# 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.
+# ==============================================================================
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import googletest
+
+
+class FloatOpsTest(xla_test.XLATestCase):
+
+  def test_float_ops(self):
+    with self.session() as session:
+      for dtype in self.float_types:
+        x = np.arange(-0.90, 0.90, 0.25)
+        self.assert_op_output_matches_expected(
+            math_ops.acos,
+            x.astype(dtype),
+            expected=np.arccos(x).astype(dtype),
+            local_session=session,
+        )
+        self.assert_op_output_matches_expected(
+            math_ops.asin,
+            x.astype(dtype),
+            expected=np.arcsin(x).astype(dtype),
+            local_session=session,
+        )
+        x = np.arange(-3, 3).reshape(1, 3, 2)
+        self.assert_op_output_matches_expected(
+            math_ops.atan,
+            x.astype(dtype),
+            expected=np.arctan(x).astype(dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.acosh,
+            np.array([1, 2, 3, 4], dtype=dtype),
+            expected=np.array(
+                [0, 1.3169579, 1.76274717, 2.06343707], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.asinh,
+            np.array([1, 2, 3, 4], dtype=dtype),
+            expected=np.array(
+                [0.88137359, 1.44363548, 1.81844646, 2.09471255], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.atanh,
+            np.array([0.1, 0.2, 0.3, 0.4], dtype=dtype),
+            expected=np.array(
+                [0.10033535, 0.20273255, 0.3095196, 0.42364893], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.ceil,
+            np.array([[-1.7, 1.2]], dtype=dtype),
+            expected=np.array([[-1, 2]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.cosh,
+            np.array([1, 2, 3, 4], dtype=dtype),
+            expected=np.array(
+                [1.54308063, 3.76219569, 10.067662, 27.30823284], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        # Disable float16 testing for now
+        if dtype != np.float16:
+          x = np.arange(-10, 10, 1).astype(dtype)
+          erf_x = session.run(math_ops.erf(x))
+          erfc_x = session.run(math_ops.erfc(x))
+
+          self.assert_op_output_matches_expected(
+              math_ops.erf,
+              x,
+              expected=erf_x,
+              local_session=session,
+          )
+          self.assert_op_output_matches_expected(
+              math_ops.erfc,
+              x,
+              expected=erfc_x,
+              local_session=session,
+          )
+
+        self.assert_op_output_matches_expected(
+            math_ops.exp,
+            np.array([[-1, 1]], dtype=dtype),
+            expected=np.array([[0.36787945, 2.7182817]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.expm1,
+            np.array([[-1, 1]], dtype=dtype),
+            expected=np.array([[-0.63212056, 1.71828183]], dtype=dtype),
+            local_session=session,
+            rtol=1e-5,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.floor,
+            np.array([[-1.7, 1.2]], dtype=dtype),
+            expected=np.array([[-2, 1]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.is_finite,
+            np.array(
+                [[-np.inf, -2, -1, 0, 0.5, 1, 2, np.inf, np.nan]], dtype=dtype
+            ),
+            expected=np.array([[0, 1, 1, 1, 1, 1, 1, 0, 0]], dtype=np.bool_),
+            local_session=session,
+        )
+
+        # Tests for tf.nn ops.
+        self.assert_op_output_matches_expected(
+            nn_ops.l2_loss,
+            np.array([[[]]], dtype=dtype),
+            expected=dtype(0),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.l2_loss,
+            dtype(4),
+            dtype(8),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.l2_loss,
+            np.array([[-2, 4]], dtype=dtype),
+            expected=dtype(10),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.reciprocal,
+            np.array([[1, 2]], dtype=dtype),
+            expected=np.array([[1, 0.5]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.log,
+            np.array([[1, 2]], dtype=dtype),
+            expected=np.array([[0, 0.69314718]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.sin,
+            np.array([[1, 2]], dtype=dtype),
+            expected=np.array([[0.841478, 0.909302]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.cos,
+            np.array([[1, 2]], dtype=dtype),
+            expected=np.array([[0.540297, -0.41614]], dtype=dtype),
+            local_session=session,
+        )
+
+        # Confirm that log1p will remain precise across a range of small values.
+        self.assert_op_output_matches_expected(
+            math_ops.log1p,
+            np.array(
+                [[1e-14, 1e-15, 0.6, 2] + [x * 1e-5 for x in range(1, 20)]],
+                dtype=dtype,
+            ),
+            expected=np.log1p(
+                np.array(
+                    [[1e-14, 1e-15, 0.6, 2] + [x * 1e-5 for x in range(1, 20)]],
+                    dtype=dtype,
+                )
+            ).astype(dtype),
+            local_session=session,
+            rtol=1e-15 if dtype == np.float64 else 1e-4,
+            atol=1e-15 if dtype == np.float64 else 1e-4,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.rint,
+            np.array(
+                [
+                    [-1.7, 1.2, 4.0, 0.0],
+                    [-3.5, -2.5, -1.5, -0.5],
+                    [0.5, 1.5, 2.5, 3.5],
+                ],
+                dtype=dtype,
+            ),
+            expected=np.array(
+                [[-2, 1, 4, 0], [-4, -2, -2, 0], [0, 2, 2, 4]], dtype=dtype
+            ),
+            local_session=session,
+        )
+        self.assert_op_output_matches_expected(
+            math_ops.round,
+            np.array(
+                [
+                    [-1.7, 1.2, 4.0, 0.0],
+                    [-3.5, -2.5, -1.5, -0.5],
+                    [0.5, 1.5, 2.5, 3.5],
+                ],
+                dtype=dtype,
+            ),
+            expected=np.array(
+                [[-2, 1, 4, 0], [-4, -2, -2, 0], [0, 2, 2, 4]], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.rsqrt,
+            np.array([[4, 16]], dtype=dtype),
+            expected=np.array([[0.5, 0.25]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.sigmoid,
+            np.array([[1, 1, 1, 1], [1, 2, 3, 4]], dtype=dtype),
+            expected=np.array(
+                [
+                    [0.7310586, 0.7310586, 0.7310586, 0.7310586],
+                    [0.7310586, 0.880797, 0.95257413, 0.98201376],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.sigmoid,
+            np.array([-300, -150, 0, 150, 300], dtype=dtype),
+            expected=np.array([0, 0, 0.5, 1, 1], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.sinh,
+            np.array([1, 2, 3, 4], dtype=dtype),
+            expected=np.array(
+                [1.17520119, 3.62686041, 10.01787493, 27.2899172], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.sqrt,
+            np.array([[4, 9]], dtype=dtype),
+            expected=np.array([[2, 3]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.tan,
+            np.array([1, 2, 3, 4], dtype=dtype),
+            expected=np.array(
+                [1.55740772, -2.18503986, -0.14254654, 1.15782128], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.tanh,
+            np.array(
+                [
+                    [1, 2, 3, 4],
+                    [np.inf, -np.inf, np.nan, 20],
+                    [19, -19, 22, -22],
+                ],
+                dtype=dtype,
+            ),
+            expected=np.array(
+                [
+                    [0.76159418, 0.96402758, 0.99505478, 0.99932933],
+                    [1.0, -1.0, np.nan, 1.0],
+                    [1.0, -1.0, 1.0, -1.0],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.log_softmax,
+            np.array([[1, 1, 1, 1], [1, 2, 3, 4]], dtype=dtype),
+            expected=np.array(
+                [
+                    [-1.3862944, -1.3862944, -1.3862944, -1.3862944],
+                    [-3.4401896, -2.4401896, -1.4401897, -0.44018969],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.elu,
+            np.array([[-1, 0, 1, -1e-6]], dtype=dtype),
+            expected=np.array(
+                [[-0.63212056, 0, 1, -9.999995e-07]], dtype=dtype
+            ),
+            rtol=1e-5,
+            atol=1e-6,
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.selu,
+            np.array([[-1, 0, 1, -1e-5]], dtype=dtype),
+            expected=np.array(
+                [[-1.11133074, 0.0, 1.05070099, -1.758090550379974e-05]],
+                dtype=dtype,
+            ),
+            rtol=1e-5,
+            atol=1e-6,
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.relu,
+            np.array([[-1, 1]], dtype=dtype),
+            expected=np.array([[0, 1]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.relu6,
+            np.array([[-0.05, 6.05, 5]], dtype=dtype),
+            expected=np.array([[0, 6, 5]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.leaky_relu,
+            np.array([[-2, -1, 0, 1, 2]], dtype=dtype),
+            expected=np.array([[-0.4, -0.2, 0.0, 1.0, 2.0]], dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.softmax,
+            np.array([1, 2, 3, 4], dtype=dtype),
+            expected=np.array(
+                [0.032058604, 0.087144323, 0.23688284, 0.64391428], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.softmax,
+            np.array([[1, 1, 1, 1], [1, 2, 3, 4]], dtype=dtype),
+            expected=np.array(
+                [
+                    [0.25, 0.25, 0.25, 0.25],
+                    [0.032058604, 0.087144323, 0.23688284, 0.64391428],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.softmax,
+            np.array([[[1, 1], [1, 1]], [[1, 2], [3, 4]]], dtype=dtype),
+            expected=np.array(
+                [
+                    [[0.5, 0.5], [0.5, 0.5]],
+                    [[0.26894142, 0.73105858], [0.26894142, 0.73105858]],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            nn_ops.softsign,
+            np.array([[-2, -1, 0, 1, 2]], dtype=dtype),
+            expected=np.array(
+                [[-0.66666669, -0.5, 0, 0.5, 0.66666669]], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.sign,
+            np.array(
+                [[-2.0, -1.0, -0.0, +0.0, 1.0, 2.0, float("nan")]], dtype=dtype
+            ),
+            expected=np.array(
+                [[-1.0, -1.0, -0.0, +0.0, 1.0, 1.0, float("nan")]], dtype=dtype
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.is_finite,
+            np.array(
+                [[42, float("inf"), -123], [float("nan"), 0, -0.0]], dtype=dtype
+            ),
+            expected=np.array(
+                [[True, False, True], [False, True, True]], dtype=np.bool_
+            ),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.lgamma,
+            np.array(0.5, dtype=dtype),
+            expected=np.array(np.log(np.pi) / 2, dtype=dtype),
+            local_session=session,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.lgamma,
+            np.array(
+                [
+                    [1, 2, 3],
+                    [4, 5, 6],
+                    [1 / 2, 3 / 2, 5 / 2],
+                    [-3 / 2, -7 / 2, -11 / 2],
+                ],
+                dtype=dtype,
+            ),
+            expected=np.array(
+                [
+                    [0, 0, np.log(2.0)],
+                    [np.log(6.0), np.log(24.0), np.log(120)],
+                    [
+                        np.log(np.pi) / 2,
+                        np.log(np.pi) / 2 - np.log(2),
+                        np.log(np.pi) / 2 - np.log(4) + np.log(3),
+                    ],
+                    [
+                        np.log(np.pi) / 2 - np.log(3) + np.log(4),
+                        np.log(np.pi) / 2 - np.log(105) + np.log(16),
+                        np.log(np.pi) / 2 - np.log(10395) + np.log(64),
+                    ],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+        # The actual result is complex. Take the real part.
+        self.assert_op_output_matches_expected(
+            math_ops.lgamma,
+            np.array([-1 / 2, -5 / 2, -9 / 2], dtype=dtype),
+            expected=np.array(
+                [
+                    np.log(np.pi) / 2 + np.log(2),
+                    np.log(np.pi) / 2 - np.log(15) + np.log(8),
+                    np.log(np.pi) / 2 - np.log(945) + np.log(32),
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+            atol=1e-4,
+        )
+
+        self.assert_op_output_matches_expected(
+            math_ops.digamma,
+            np.array(
+                [
+                    [1.0, 0.5, 1 / 3.0],
+                    [0.25, 1 / 6.0, 0.125],
+                    [2.0, 3.0, 4.0],
+                    [6.0, 8.0, 9.0],
+                ],
+                dtype=dtype,
+            ),
+            expected=np.array(
+                [
+                    [
+                        -np.euler_gamma,
+                        -2 * np.log(2) - np.euler_gamma,
+                        -np.pi / 2 / np.sqrt(3)
+                        - 3 * np.log(3) / 2
+                        - np.euler_gamma,
+                    ],
+                    [
+                        -np.pi / 2 - 3 * np.log(2) - np.euler_gamma,
+                        -np.pi * np.sqrt(3) / 2
+                        - 2 * np.log(2)
+                        - 3 * np.log(3) / 2
+                        - np.euler_gamma,
+                        -np.pi / 2
+                        - 4 * np.log(2)
+                        - (
+                            np.pi
+                            + np.log(2 + np.sqrt(2))
+                            - np.log(2 - np.sqrt(2))
+                        )
+                        / np.sqrt(2)
+                        - np.euler_gamma,
+                    ],
+                    [
+                        1 - np.euler_gamma,
+                        1.5 - np.euler_gamma,
+                        11 / 6.0 - np.euler_gamma,
+                    ],
+                    [
+                        137 / 60.0 - np.euler_gamma,
+                        363 / 140.0 - np.euler_gamma,
+                        761 / 280.0 - np.euler_gamma,
+                    ],
+                ],
+                dtype=dtype,
+            ),
+            local_session=session,
+        )
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/ftrl_ops_test.py

@@ -0,0 +1,198 @@
+# 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.
+# ==============================================================================
+"""Test cases for ftrl ("follow the regularized leader") operations."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gen_training_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.platform import googletest
+
+
+class ResourceApplyFtrlTest(xla_test.XLATestCase):
+  """Test cases for ftrl ops."""
+
+  def setUp(self):
+    super().setUp()
+    self.rewrite_ops_for_tpu = ("TPU" in self.device and
+                                test_util.is_mlir_bridge_enabled())
+
+  def _eval(self, var, accum, linear, grad, lr, l1, l2, l2_shrinkage=0,
+            lr_power=1, multiply_linear_by_lr=False):
+    dtype = np.float32
+    var = np.array(var, dtype=dtype)
+    accum = np.array(accum, dtype=dtype)
+    linear = np.array(linear, dtype=dtype)
+    grad = np.array(grad, dtype=dtype)
+    use_v2 = bool(l2_shrinkage)
+    with self.session() as session:
+      with self.test_scope():
+        lr = constant_op.constant(lr, dtype=dtype)
+        l1 = constant_op.constant(l1, dtype=dtype)
+        l2 = constant_op.constant(l2, dtype=dtype)
+        l2_shrinkage = constant_op.constant(l2_shrinkage, dtype=dtype)
+        lr_power = constant_op.constant(lr_power, dtype=dtype)
+        v_var = resource_variable_ops.ResourceVariable(var, dtype=dtype)
+        v_accum = resource_variable_ops.ResourceVariable(accum, dtype=dtype)
+        v_linear = resource_variable_ops.ResourceVariable(linear, dtype=dtype)
+        session.run(v_var.create)
+        session.run(v_accum.create)
+        session.run(v_linear.create)
+        assert not (use_v2 and multiply_linear_by_lr)
+        if use_v2:
+          session.run(gen_training_ops.resource_apply_ftrl_v2(
+              v_var.handle, v_accum.handle, v_linear.handle,
+              grad, lr, l1, l2, l2_shrinkage, lr_power,
+              multiply_linear_by_lr=multiply_linear_by_lr))
+        else:
+          session.run(gen_training_ops.resource_apply_ftrl(
+              v_var.handle, v_accum.handle, v_linear.handle,
+              grad, lr, l1, l2, lr_power,
+              multiply_linear_by_lr=multiply_linear_by_lr))
+        return (v_var.read_value().eval().reshape(var.shape),
+                v_accum.read_value().eval().reshape(accum.shape),
+                v_linear.read_value().eval().reshape(linear.shape))
+
+  def testAccum(self):
+    """Test that accum is updated with grad^2."""
+    accum = np.array([[[1, 3], [2, 5], [6, 8]]])
+    grad = np.array([[[1, 3], [2, 5], [6, 8]]])
+    _, new_accum, _ = self._eval(
+        var=np.zeros((1, 3, 2)),
+        accum=accum,
+        linear=np.zeros((1, 3, 2)),
+        grad=grad,
+        lr=7, l1=3, l2=7, lr_power=2)
+    self.assertAllClose(accum + grad*grad, new_accum)
+
+  def testLinearNoGradient(self):
+    """Test that if accum_new == accum, linear doesn't change."""
+    _, _, linear = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=[[[1, 3], [2, 5], [6, 8]]],
+        linear=[[[1, 2], [3, 4], [5, 6]]],
+        grad=np.zeros((1, 3, 2)),  # make accum_new == acum
+        lr=1, l1=3, l2=7, lr_power=2)
+    self.assertAllClose([[[1, 2], [3, 4], [5, 6]]], linear)
+
+  def testLinear(self):
+    """Test the linear update for new_linear=2 and linear=1."""
+    _, _, linear = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=1, l1=3, l2=7, lr_power=2)
+    self.assertAllClose(1.75 * np.ones((1, 3, 2)), linear)
+
+  def testLR(self):
+    """Test that the linear update is divided by lr."""
+    _, _, linear = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=5, l1=3, l2=7, lr_power=-1)
+    self.assertAllClose(0.8 * np.ones((1, 3, 2)), linear)
+
+  def testVar(self):
+    """Test computation of var with linear=1.5, quadratic=1."""
+    var, _, _ = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=1, l1=1, l2=0.25, lr_power=1)
+    self.assertAllClose(-0.5 * np.ones((1, 3, 2)), var)
+
+  def testVarClipped(self):
+    """Test that var becomes 0 if |linear| < l1."""
+    var, _, _ = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=1, l1=1.6, l2=0.25, lr_power=1)
+    self.assertAllClose(np.zeros((1, 3, 2)), var)
+
+  def testQuadratic(self):
+    """Test that quadratic (here: -2) is the divisor of var."""
+    var, _, _ = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=1, l1=1, l2=-1.25, lr_power=1)
+    self.assertAllClose(0.25 * np.ones((1, 3, 2)), var)
+
+  def testL2Shrinkage(self):
+    """Test that 2 * l2_shrinkage * var is *not* added to the gradient."""
+    _, accum, _ = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.zeros((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.zeros((1, 3, 2)),
+        lr=7, l1=3, l2=7, lr_power=2, l2_shrinkage=0.5)
+    self.assertAllClose(np.zeros((1, 3, 2)), accum)
+
+  def testL2ShrinkageOnLinear(self):
+    """Test that 2 * l2_shrinkage * var is added to linear."""
+    _, _, linear = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.zeros((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.zeros((1, 3, 2)),
+        lr=2, l1=3, l2=7, lr_power=0, l2_shrinkage=11)
+    self.assertAllClose(22 * np.ones((1, 3, 2)), linear)
+
+  def testMultiplyLinearByLR(self):
+    """Test multiply_linear_by_lr = true for the linear variable."""
+    _, _, linear = self._eval(
+        var=np.zeros((1, 3, 2)),
+        accum=np.zeros((1, 3, 2)),
+        linear=np.ones((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=6, l1=1, l2=-1.25, lr_power=0,
+        multiply_linear_by_lr=True)
+    self.assertAllClose(7 * np.ones((1, 3, 2)), linear)
+
+  def testMultiplyLinearByLRClipping(self):
+    """Test that multiply_linear_by_lr = true scales the clip margins."""
+    var, _, _ = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=3, l1=1.0, l2=0.25, lr_power=1,
+        multiply_linear_by_lr=True)
+    self.assertAllClose(-0.25 * np.ones((1, 3, 2)), var)
+
+  def testMultiplyLinearByLRClipZero(self):
+    """Test that multiply_linear_by_lr = true still clips to 0."""
+    var, _, _ = self._eval(
+        var=np.ones((1, 3, 2)),
+        accum=np.ones((1, 3, 2)),
+        linear=np.zeros((1, 3, 2)),
+        grad=np.ones((1, 3, 2)),
+        lr=3, l1=1.2, l2=0.25, lr_power=1,
+        multiply_linear_by_lr=True)
+    self.assertAllClose(np.zeros((1, 3, 2)), var)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/ftrl_test.py

@@ -0,0 +1,352 @@
+# 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.
+# ==============================================================================
+"""Tests for Ftrl optimizer."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+from tensorflow.python.training import ftrl
+from tensorflow.python.training import gradient_descent
+
+class FtrlOptimizerTest(xla_test.XLATestCase):
+
+  def initVariableAndGradient(self, dtype):
+    var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+    var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+    grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+    grads1 = constant_op.constant([0.02, 0.04], dtype=dtype)
+
+    return var0, var1, grads0, grads1
+
+  def equivAdagradTest_FtrlPart(self, steps, dtype):
+    var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+    opt = ftrl.FtrlOptimizer(
+        3.0,
+        learning_rate_power=-0.5,  # using Adagrad learning rate
+        initial_accumulator_value=0.1,
+        l1_regularization_strength=0.0,
+        l2_regularization_strength=0.0)
+    ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+    # Fetch params to validate initial values
+    self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+    self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+    # Run Ftrl for a few steps
+    for _ in range(steps):
+      ftrl_update.run()
+
+    return self.evaluate(var0), self.evaluate(var1)
+
+  def equivAdagradTest_AdagradPart(self, steps, dtype):
+    var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+    opt = adagrad.AdagradOptimizer(3.0, initial_accumulator_value=0.1)
+    adagrad_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+    # Fetch params to validate initial values
+    self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+    self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+    # Run Adagrad for a few steps
+    for _ in range(steps):
+      adagrad_update.run()
+
+    return self.evaluate(var0), self.evaluate(var1)
+
+  def equivGradientDescentTest_FtrlPart(self, steps, dtype):
+    var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+    opt = ftrl.FtrlOptimizer(
+        3.0,
+        learning_rate_power=-0.0,  # using Fixed learning rate
+        initial_accumulator_value=0.1,
+        l1_regularization_strength=0.0,
+        l2_regularization_strength=0.0)
+    ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+    # Fetch params to validate initial values
+    self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+    self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+    # Run Ftrl for a few steps
+    for _ in range(steps):
+      ftrl_update.run()
+
+    return self.evaluate(var0), self.evaluate(var1)
+
+  def equivGradientDescentTest_GradientDescentPart(self, steps, dtype):
+    var0, var1, grads0, grads1 = self.initVariableAndGradient(dtype)
+    opt = gradient_descent.GradientDescentOptimizer(3.0, name="sgd")
+    sgd_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+    # Fetch params to validate initial values
+    self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+    self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+    # Run GradientDescent for a few steps
+    for _ in range(steps):
+      sgd_update.run()
+
+    return self.evaluate(var0), self.evaluate(var1)
+
+  def testFtrlwithoutRegularization(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([0.0, 0.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.0,
+            l2_regularization_strength=0.0)
+        ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+        self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+        # Run 3 steps FTRL
+        for _ in range(3):
+          ftrl_update.run()
+
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-2.60260963, -4.29698515]),
+            self.evaluate(var0),
+            float_rtol=1e-4,
+            half_rtol=1e-2)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.28432083, -0.56694895]),
+            self.evaluate(var1),
+            float_rtol=1e-5,
+            half_rtol=1e-2)
+
+  def testFtrlwithoutRegularization2(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.0,
+            l2_regularization_strength=0.0)
+        ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+        # Run 3 steps FTRL
+        for _ in range(3):
+          ftrl_update.run()
+
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-2.55607247, -3.98729396]),
+            self.evaluate(var0),
+            1e-5,
+            1e-5,
+            float_rtol=1e-4)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.28232238, -0.56096673]), self.evaluate(var1), 1e-5,
+            1e-5)
+
+  def testFtrlWithL1(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=0.0)
+        ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+        # Run 10 steps FTRL
+        for _ in range(10):
+          ftrl_update.run()
+
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-7.66718769, -10.91273689]),
+            self.evaluate(var0),
+            rtol=1e-4,
+            bfloat16_rtol=1e-1,
+            bfloat16_atol=1e-1)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.93460727, -1.86147261]),
+            self.evaluate(var1),
+            rtol=1e-4)
+
+  def testFtrlWithL1_L2(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=2.0)
+        ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+        # Run 10 steps FTRL
+        for _ in range(10):
+          ftrl_update.run()
+
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-0.24059935, -0.46829352]),
+            self.evaluate(var0),
+            rtol=1e-5)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.02406147, -0.04830509]),
+            self.evaluate(var1),
+            rtol=1e-5)
+
+  def testFtrlWithL1_L2_L2Shrinkage(self):
+    """Test the new FTRL op with support for l2 shrinkage.
+
+    The addition of this parameter which places a constant pressure on weights
+    towards the origin causes the gradient descent trajectory to differ. The
+    weights will tend to have smaller magnitudes with this parameter set.
+    """
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([4.0, 3.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.02], dtype=dtype)
+        opt = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=2.0,
+            l2_shrinkage_regularization_strength=0.1)
+        ftrl_update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Fetch params to validate initial values
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([4.0, 3.0], self.evaluate(var1))
+
+        # Run 10 steps FTRL
+        for _ in range(10):
+          ftrl_update.run()
+
+        # Validate updated params
+        self.assertAllCloseAccordingToType(
+            np.array([-0.22578996, -0.44345799]),
+            self.evaluate(var0),
+            rtol=1e-4)
+        self.assertAllCloseAccordingToType(
+            np.array([-0.14378493, -0.13229476]),
+            self.evaluate(var1),
+            rtol=1e-4)
+
+  def testFtrlWithL2ShrinkageDoesNotChangeLrSchedule(self):
+    """Verifies that l2 shrinkage in FTRL does not change lr schedule."""
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.2], dtype=dtype)
+        grads1 = constant_op.constant([0.1, 0.2], dtype=dtype)
+
+        opt0 = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=2.0,
+            l2_shrinkage_regularization_strength=0.1)
+        opt1 = ftrl.FtrlOptimizer(
+            3.0,
+            initial_accumulator_value=0.1,
+            l1_regularization_strength=0.001,
+            l2_regularization_strength=2.0)
+        update0 = opt0.apply_gradients([(grads0, var0)])
+        update1 = opt1.apply_gradients([(grads1, var1)])
+        self.evaluate(variables.global_variables_initializer())
+
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var0))
+        self.assertAllCloseAccordingToType([1.0, 2.0], self.evaluate(var1))
+
+        # Run 10 steps FTRL
+        for _ in range(10):
+          update0.run()
+          update1.run()
+
+        # var0 is experiencing L2 shrinkage so it should be smaller than var1
+        # in magnitude.
+        self.assertTrue((var0.eval()**2 < self.evaluate(var1)**2).all())
+        accum0 = list(opt0._slots["accum"].values())[0].eval()
+        accum1 = list(opt1._slots["accum"].values())[0].eval()
+        # L2 shrinkage should not change how we update grad accumulator.
+        self.assertAllCloseAccordingToType(accum0, accum1)
+
+  # When variables are initialized with Zero, FTRL-Proximal has two properties:
+  # 1. Without L1&L2 but with fixed learning rate, FTRL-Proximal is identical
+  # with GradientDescent.
+  # 2. Without L1&L2 but with adaptive learning rate, FTRL-Proximal is identical
+  # with Adagrad.
+  # So, basing on these two properties, we test if our implementation of
+  # FTRL-Proximal performs same updates as Adagrad or GradientDescent.
+  def testEquivAdagradwithoutRegularization(self):
+    steps = 5
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        val0, val1 = self.equivAdagradTest_FtrlPart(steps, dtype)
+      with self.session(), self.test_scope():
+        val2, val3 = self.equivAdagradTest_AdagradPart(steps, dtype)
+
+    self.assertAllCloseAccordingToType(val0, val2, rtol=1e-4, half_rtol=1e-2)
+    self.assertAllCloseAccordingToType(val1, val3, rtol=1e-4, half_rtol=1e-2)
+
+  def testEquivGradientDescentwithoutRegularization(self):
+    steps = 5
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        val0, val1 = self.equivGradientDescentTest_FtrlPart(steps, dtype)
+      with self.session(), self.test_scope():
+        val2, val3 = self.equivGradientDescentTest_GradientDescentPart(
+            steps, dtype)
+
+    self.assertAllCloseAccordingToType(val0, val2, rtol=1e-5)
+    self.assertAllCloseAccordingToType(val1, val3, rtol=1e-5)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/function_test.py

@@ -0,0 +1,149 @@
+# 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.
+# ==============================================================================
+"""Test cases for Tensorflow functions."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import googletest
+
+
+class FunctionTest(xla_test.XLATestCase):
+
+  def testFunction(self):
+    """Executes a simple TensorFlow function."""
+
+    def APlus2B(a, b):
+      return a + b * 2
+
+    aval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32)
+    bval = np.array([5, 6, 7, 8]).reshape([2, 2]).astype(np.float32)
+    expected = APlus2B(aval, bval)
+
+    with self.session():
+
+      @function.Defun(dtypes.float32, dtypes.float32)
+      def Foo(a, b):
+        return APlus2B(a, b)
+
+      a = constant_op.constant(aval, name="a")
+      b = constant_op.constant(bval, name="b")
+      with self.test_scope():
+        call_f = Foo(a, b)
+      result = self.evaluate(call_f)
+    self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testNestedFunctions(self):
+    """Executes two nested TensorFlow functions."""
+
+    def TimesTwo(x):
+      return x * 2
+
+    def APlus2B(a, b):
+      return a + TimesTwo(b)
+
+    aval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32)
+    bval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32)
+    expected = APlus2B(aval, bval)
+
+    with self.session():
+
+      @function.Defun(dtypes.float32, dtypes.float32)
+      def Foo(a, b):
+        return APlus2B(a, b)
+
+      a = constant_op.constant(aval, name="a")
+      b = constant_op.constant(bval, name="b")
+      with self.test_scope():
+        call_g = Foo(a, b)
+      result = self.evaluate(call_g)
+    self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testFunctionMultipleRetvals(self):
+    """Executes a function with multiple return values."""
+
+    # This function will run on the XLA device
+    def Func(a, b):
+      return a + b, a - b
+
+    aval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32)
+    bval = np.array([5, 6, 7, 8]).reshape([2, 2]).astype(np.float32)
+    expected = Func(aval, bval)
+
+    with self.session():
+
+      @function.Defun(dtypes.float32, dtypes.float32)
+      def Foo(a, b):
+        return Func(a, b)
+
+      a = constant_op.constant(aval, name="a")
+      b = constant_op.constant(bval, name="b")
+      with self.test_scope():
+        call_f = Foo(a, b)
+      result = self.evaluate(call_f)
+    self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testCompileTimeConstantsInDefun(self):
+    """Tests that XLA handles compile-time constants in defuns."""
+    with self.session() as sess:
+
+      @function.Defun(dtypes.float32, dtypes.int32, dtypes.int32)
+      def Foo(a, c, d):
+        # c and d must be known at compile time
+        x = array_ops.slice(a, c, d)
+        return x
+
+      a = array_ops.placeholder(dtypes.float32)
+      c = array_ops.placeholder(dtypes.int32, shape=[4])
+      d = array_ops.placeholder(dtypes.int32, shape=[4])
+      with self.test_scope():
+        call_f = Foo(a, c, d)
+      result = sess.run(call_f, feed_dict={
+          a: np.ones([1, 4, 4, 1]),
+          c: [0, 0, 0, 0],
+          d: [1, 2, 2, 1]})
+
+    self.assertAllEqual(np.ones([1, 2, 2, 1]), result)
+
+  # TODO(b/36139787): Re-enable this test when noinline works again.
+  def DISABLED_testFunctionsNoInline(self):
+
+    @function.Defun(dtypes.float32, noinline=True)
+    def TimesTwo(x):
+      return x * 2
+
+    @function.Defun(dtypes.float32, dtypes.float32)
+    def APlus2B(a, b):
+      return a + TimesTwo(b)
+
+    aval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32)
+    bval = np.array([4, 3, 2, 1]).reshape([2, 2]).astype(np.float32)
+    expected = aval + bval * 2
+
+    with self.session() as sess:
+      with self.test_scope():
+        a = array_ops.placeholder(dtypes.float32, name="a")
+        b = array_ops.placeholder(dtypes.float32, name="b")
+        call = APlus2B(a, b)
+      result = sess.run(call, {a: aval, b: bval})
+    self.assertAllClose(result, expected, rtol=1e-3)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/fused_batchnorm_test.py

@@ -0,0 +1,355 @@
+# 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 fused batch norm operations."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import test_utils
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import test
+
+DATA_FORMATS = (
+    ("_data_format_NHWC", "NHWC"),
+    ("_data_format_NCHW", "NCHW"),
+)
+
+DATA_FORMATS_AND_AVG_FACTORS = (
+    ("_data_format_NHWC_no_averaging", "NHWC", 1.0),
+    ("_data_format_NHWC_averaging", "NHWC", 0.6),
+    ("_data_format_NCHW_no_averaging", "NCHW", 1.0),
+    ("_data_format_NCHW_averaging", "NCHW", 0.6),
+)
+
+
+class FusedBatchNormTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _reference_training(self, x, scale, offset, old_mean, old_var, epsilon,
+                          exponential_avg_factor, data_format):
+    if data_format != "NHWC":
+      raise ValueError("data_format must be NHWC, got %s." % data_format)
+    x_square = x * x
+    x_square_sum = np.sum(x_square, (0, 1, 2))
+    x_sum = np.sum(x, axis=(0, 1, 2))
+    element_count = np.size(x) / int(np.shape(x)[-1])
+    mean = x_sum / element_count
+    var = x_square_sum / element_count - mean * mean
+    factor = element_count / max(element_count - 1, 1)
+    corrected_var = var * factor
+    normalized = (x - mean) / np.sqrt(var + epsilon)
+    if exponential_avg_factor != 1.0:
+      mean = (1.0 -
+              exponential_avg_factor) * old_mean + exponential_avg_factor * mean
+      corrected_var = (1.0 - exponential_avg_factor
+                      ) * old_var + exponential_avg_factor * corrected_var
+    return (normalized * scale + offset), mean, var, corrected_var
+
+  def _reference_grad(self, x, grad_y, scale, mean, var, epsilon, data_format):
+    # Use the following formulas to calculate gradients:
+    # grad_scale =
+    #   sum(grad_y * (x - mean)) * rsqrt(var + epsilon)
+    #
+    # grad_offset = sum(output_y)
+    #
+    # grad_x =
+    #   1/N * scale * rsqrt(var + epsilon) * (N * grad_y - sum(grad_y) -
+    #   (x - mean) * sum(grad_y * (x - mean)) / (var + epsilon))
+    if data_format != "NHWC":
+      raise ValueError("data_format must be NHWC, got %s." % data_format)
+    grad_x = scale * (grad_y - np.mean(grad_y, axis=(0, 1, 2)) -
+                      (x - mean) * np.mean(grad_y *
+                                           (x - mean), axis=(0, 1, 2)) /
+                      (var + epsilon)) / np.sqrt(var + epsilon)
+    grad_scale = np.sum(
+        grad_y * (x - mean) / np.sqrt(var + epsilon), axis=(0, 1, 2))
+    grad_offset = np.sum(grad_y, axis=(0, 1, 2))
+    return grad_x, grad_scale, grad_offset
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testInference(self, data_format):
+    channel = 3
+    x_shape = [2, 2, 6, channel]
+    scale_shape = [channel]
+    x_val = np.random.random_sample(x_shape).astype(np.float32)
+    scale_val = np.random.random_sample(scale_shape).astype(np.float32)
+    offset_val = np.random.random_sample(scale_shape).astype(np.float32)
+    epsilon = 0.001
+    exponential_avg_factor = 1.0
+    data_format_src = "NHWC"
+    y_ref, mean_ref, var_ref, _ = self._reference_training(
+        x_val, scale_val, offset_val, None, None, epsilon,
+        exponential_avg_factor, data_format_src)
+
+    with self.session() as sess, self.test_scope():
+      # To avoid constant folding
+      x_val_converted = test_utils.ConvertBetweenDataFormats(
+          x_val, data_format_src, data_format)
+      y_ref_converted = test_utils.ConvertBetweenDataFormats(
+          y_ref, data_format_src, data_format)
+
+      t_val = array_ops.placeholder(
+          np.float32, shape=x_val_converted.shape, name="x")
+      scale = array_ops.placeholder(np.float32, shape=scale_shape, name="scale")
+      offset = array_ops.placeholder(
+          np.float32, shape=scale_shape, name="offset")
+      y, mean, variance = nn.fused_batch_norm(
+          t_val,
+          scale,
+          offset,
+          mean=mean_ref,
+          variance=var_ref,
+          epsilon=epsilon,
+          data_format=data_format,
+          is_training=False)
+
+      y_val, _, _ = sess.run([y, mean, variance], {
+          t_val: x_val_converted,
+          scale: scale_val,
+          offset: offset_val
+      })
+      self.assertAllClose(y_val, y_ref_converted, atol=1e-3)
+
+  def _testLearning(self, use_gradient_checker, data_format,
+                    exponential_avg_factor):
+    channel = 3
+    x_shape = [2, 2, 6, channel]
+    scale_shape = [channel]
+    x_val = np.random.random_sample(x_shape).astype(np.float32)
+    scale_val = np.random.random_sample(scale_shape).astype(np.float32)
+    offset_val = np.random.random_sample(scale_shape).astype(np.float32)
+    mean_val = np.random.random_sample(scale_shape).astype(np.float32)
+    var_val_corr = np.random.random_sample(scale_shape).astype(np.float32)
+    epsilon = 0.001
+    data_format_src = "NHWC"
+    # When in training mode, fused_batchnorm applies an implicit Bessel's
+    # correction. So we have to use the corrected variance here, as well.
+    y_ref, mean_ref, _, var_ref_corr = self._reference_training(
+        x_val, scale_val, offset_val, mean_val, var_val_corr, epsilon,
+        exponential_avg_factor, data_format_src)
+
+    with self.session() as sess, self.test_scope():
+      # To avoid constant folding
+      x_val_converted = test_utils.ConvertBetweenDataFormats(
+          x_val, data_format_src, data_format)
+      y_ref_converted = test_utils.ConvertBetweenDataFormats(
+          y_ref, data_format_src, data_format)
+
+      t_val = array_ops.placeholder(
+          np.float32, shape=x_val_converted.shape, name="x")
+      scale = array_ops.placeholder(np.float32, shape=scale_shape, name="scale")
+      offset = array_ops.placeholder(
+          np.float32, shape=scale_shape, name="offset")
+      if exponential_avg_factor == 1.0:
+        old_mean = None
+        old_var = None
+      else:
+        old_mean = array_ops.placeholder(
+            np.float32, shape=scale_shape, name="old_mean")
+        old_var = array_ops.placeholder(
+            np.float32, shape=scale_shape, name="old_var")
+      y, mean, var = nn.fused_batch_norm(
+          t_val,
+          scale,
+          offset,
+          mean=old_mean,
+          variance=old_var,
+          epsilon=epsilon,
+          exponential_avg_factor=exponential_avg_factor,
+          data_format=data_format,
+          is_training=True)
+      if exponential_avg_factor == 1.0:
+        feed_dict = {
+            t_val: x_val_converted,
+            scale: scale_val,
+            offset: offset_val,
+        }
+      else:
+        feed_dict = {
+            t_val: x_val_converted,
+            scale: scale_val,
+            offset: offset_val,
+            old_mean: mean_val,
+            old_var: var_val_corr
+        }
+      # Check gradient.
+      if use_gradient_checker:
+        err = gradient_checker.compute_gradient_error(
+            t_val,
+            x_val_converted.shape,
+            y,
+            x_val_converted.shape,
+            extra_feed_dict=feed_dict)
+        self.assertLess(err, 1e-3)
+
+      y_tf, mean_tf, var_tf = sess.run([y, mean, var], feed_dict)
+      self.assertAllClose(y_tf, y_ref_converted, atol=1e-3)
+      self.assertAllClose(mean_tf, mean_ref, atol=1e-3)
+      self.assertAllClose(var_tf, var_ref_corr, atol=1e-3)
+
+  @parameterized.named_parameters(*DATA_FORMATS_AND_AVG_FACTORS)
+  def testLearning(self, data_format, exponential_avg_factor):
+    self._testLearning(False, data_format, exponential_avg_factor)
+
+  @parameterized.named_parameters(*DATA_FORMATS_AND_AVG_FACTORS)
+  def testLearningWithGradientChecker(self, data_format,
+                                      exponential_avg_factor):
+    self._testLearning(True, data_format, exponential_avg_factor)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testGradientTraining(self, data_format):
+    # disable_mlir_bridge for GPUs as there is no legalization for GPU with
+    # MLIR.
+    # TODO(b/189039456): Customize FusedBatchNorm legalization for GPU in MLIR.
+    if test_util.is_mlir_bridge_enabled() and self.device == "XLA_GPU":
+      self.skipTest("b/189039456")
+
+    # TODO(b/64270657): Use gradient_checker here in addition to comparing with
+    # this reference implementation.
+    channel = 3
+    x_shape = [2, 2, 6, channel]
+    scale_shape = [channel]
+    grad_val = np.random.random_sample(x_shape).astype(np.float32)
+    x_val = np.random.random_sample(x_shape).astype(np.float32)
+    scale_val = np.random.random_sample(scale_shape).astype(np.float32)
+    mean_val = np.random.random_sample(scale_shape).astype(np.float32)
+    var_val = np.random.random_sample(scale_shape).astype(np.float32)
+    epsilon = 0.001
+
+    # The TensorFlow FusedBatchNormGrad training operation takes two inputs with
+    # implementation defined values.  In theory the only correct value these
+    # inputs are the corresponding reserve_space_{1|2} outputs from the
+    # FusedBatchNorm training operation.  However, in practice, we rely on the
+    # first one being mean on {C|G}PU, and the second one being variance on CPU
+    # and inverse(sqrt(variance + epsilon)) on GPU (we test this assumption
+    # separately).
+    reserve_space_1_val = mean_val
+    if self.device == "XLA_GPU":
+      reserve_space_2_val = np.reciprocal(np.sqrt(var_val + epsilon))
+    else:
+      reserve_space_2_val = var_val
+
+    data_format_src = "NHWC"
+    grad_x_ref, grad_scale_ref, grad_offset_ref = self._reference_grad(
+        x_val, grad_val, scale_val, mean_val, var_val, epsilon, data_format_src)
+
+    with self.session() as sess, self.test_scope():
+      grad_val_converted = test_utils.ConvertBetweenDataFormats(
+          grad_val, data_format_src, data_format)
+      x_val_converted = test_utils.ConvertBetweenDataFormats(
+          x_val, data_format_src, data_format)
+      grad_x_ref_converted = test_utils.ConvertBetweenDataFormats(
+          grad_x_ref, data_format_src, data_format)
+
+      grad = array_ops.placeholder(
+          np.float32, shape=x_val_converted.shape, name="grad")
+      x = array_ops.placeholder(
+          np.float32, shape=x_val_converted.shape, name="x")
+      reserve_space_1 = array_ops.placeholder(
+          np.float32, shape=scale_shape, name="reserve_space_1")
+      reserve_space_2 = array_ops.placeholder(
+          np.float32, shape=scale_shape, name="reserve_space_2")
+      scale = array_ops.placeholder(np.float32, shape=scale_shape, name="scale")
+      grad_x, grad_scale, grad_offset, _, _ = gen_nn_ops.fused_batch_norm_grad(
+          grad,
+          x,
+          scale,
+          reserve_space_1,
+          reserve_space_2,
+          data_format=data_format,
+          is_training=True)
+
+      grad_x_val, grad_scale_val, grad_offset_val = sess.run(
+          [grad_x, grad_scale, grad_offset], {
+              grad: grad_val_converted,
+              x: x_val_converted,
+              reserve_space_1: reserve_space_1_val,
+              reserve_space_2: reserve_space_2_val,
+              scale: scale_val
+          })
+
+      self.assertAllClose(grad_x_val, grad_x_ref_converted, atol=1e-2)
+      self.assertAllClose(grad_scale_val, grad_scale_ref, atol=1e-2)
+      self.assertAllClose(grad_offset_val, grad_offset_ref, atol=1e-3)
+
+  @parameterized.named_parameters(*DATA_FORMATS)
+  def testGradientInference(self, data_format):
+    # TODO(b/64270657): Use gradient_checker here in addition to comparing with
+    # this reference implementation.
+    channel = 3
+    x_shape = [2, 2, 6, channel]
+    scale_shape = [channel]
+    grad_val = np.random.random_sample(x_shape).astype(np.float32)
+    x_val = np.random.random_sample(x_shape).astype(np.float32)
+    scale_val = np.random.random_sample(scale_shape).astype(np.float32)
+    mean_val = np.random.random_sample(scale_shape).astype(np.float32)
+    var_val = np.random.random_sample(scale_shape).astype(np.float32)
+    data_format_src = "NHWC"
+
+    with self.session() as sess, self.test_scope():
+      grad_val_converted = test_utils.ConvertBetweenDataFormats(
+          grad_val, data_format_src, data_format)
+      x_val_converted = test_utils.ConvertBetweenDataFormats(
+          x_val, data_format_src, data_format)
+
+      grad = array_ops.placeholder(
+          np.float32, shape=x_val_converted.shape, name="grad")
+      x = array_ops.placeholder(
+          np.float32, shape=x_val_converted.shape, name="x")
+      mean = array_ops.placeholder(np.float32, shape=scale_shape, name="mean")
+      var = array_ops.placeholder(np.float32, shape=scale_shape, name="var")
+      scale = array_ops.placeholder(np.float32, shape=scale_shape, name="scale")
+      with self.test_scope():
+        out = gen_nn_ops.fused_batch_norm_grad(
+            grad,
+            x,
+            scale,
+            mean,
+            var,
+            data_format=data_format,
+            is_training=False)
+        grad_x, grad_scale, grad_offset, _, _ = out
+
+      ref_x, ref_scale, ref_offset, _, _ = gen_nn_ops.fused_batch_norm_grad(
+          grad, x, scale, mean, var, data_format=data_format, is_training=False)
+
+      grad_x_val, grad_scale_val, grad_offset_val, = sess.run(
+          [grad_x, grad_scale, grad_offset], {
+              grad: grad_val_converted,
+              x: x_val_converted,
+              mean: mean_val,
+              var: var_val,
+              scale: scale_val
+          })
+      grad_x_ref, grad_scale_ref, grad_offset_ref, = sess.run(
+          [ref_x, ref_scale, ref_offset], {
+              grad: grad_val_converted,
+              x: x_val_converted,
+              mean: mean_val,
+              var: var_val,
+              scale: scale_val
+          })
+
+      self.assertAllClose(grad_x_val, grad_x_ref, atol=1e-2)
+      self.assertAllClose(grad_scale_val, grad_scale_ref, atol=1e-2)
+      self.assertAllClose(grad_offset_val, grad_offset_ref, atol=1e-3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/gather_nd_op_test.py

@@ -0,0 +1,175 @@
+# Copyright 2018 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 tensorflow.ops.tf.gather_nd."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class GatherNdTest(xla_test.XLATestCase):
+
+  def _runGather(self, params, indices, bad_indices_policy=""):
+    with self.session():
+      paramsp = array_ops.placeholder(params.dtype)
+      indicesp = array_ops.placeholder(indices.dtype)
+      with self.test_scope():
+        gather_nd_t = array_ops.gather_nd(
+            paramsp, indicesp, bad_indices_policy=bad_indices_policy
+        )
+      feed_dict = {paramsp: params, indicesp: indices}
+      return gather_nd_t.eval(feed_dict=feed_dict)
+
+  def testSimpleDtype(self):
+    for dtype in self.numeric_types:
+      self.assertAllEqual(
+          np.array([7, 7, 8], dtype=dtype),
+          self._runGather(
+              np.array([8, 1, 2, 3, 7, 5], dtype=dtype),
+              np.array([[4], [4], [0]], np.int32)))
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testEmptyIndicesAndParamstAndEmptyParamsOk(self):
+    with self.session():
+      params = np.ones((3, 3), dtype=np.float32)
+
+      indices_empty = np.empty((0, 2), dtype=np.int32)
+      gather_nd_ok_val = self._runGather(params, indices_empty)
+      self.assertAllClose(np.empty((0,), dtype=np.float32), gather_nd_ok_val)
+
+      indices_empty = np.empty((0, 1), dtype=np.int32)
+      gather_nd_ok_val = self._runGather(params, indices_empty)
+      self.assertAllClose(np.empty((0, 3), dtype=np.float32), gather_nd_ok_val)
+
+      params_empty = np.empty((0, 3), dtype=np.float32)
+      indices_empty = np.empty((0, 2), dtype=np.int32)
+      gather_nd_ok_val = self._runGather(params_empty, indices_empty)
+      self.assertAllClose(np.empty((0,), dtype=np.float32), gather_nd_ok_val)
+
+      # Zero sized indices results in a constant of 0
+      params_empty = np.empty((0, 3), dtype=np.float32)
+      indices_nonempty = np.zeros((1, 2), dtype=np.int32)
+      gather_nd_ok_val = self._runGather(params_empty, indices_nonempty)
+      self.assertAllEqual(gather_nd_ok_val, np.zeros([3]))
+
+  def testIndexScalar(self):
+    params = np.array(
+        [[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], dtype=np.float32).T
+    indices = np.array([4, 1], dtype=np.int32)
+    gather_nd_val = self._runGather(params, indices)
+    self.assertAllEqual(np.array(7), gather_nd_val)
+
+  def testParamsRankLargerThanIndexIndexScalarSlices(self):
+    params = np.array(
+        [[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], dtype=np.float32).T
+    indices = np.array(
+        [
+            4,
+        ], dtype=np.int32)
+    gather_nd_val = self._runGather(params, indices)
+    self.assertAllEqual(np.array([-7, 7]), gather_nd_val)
+
+  def testParamsRankLargerThanIndexSlices(self):
+    params = np.array(
+        [[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]], dtype=np.float32).T
+    indices = np.array([[4], [4], [0]], np.int32)
+    gather_nd_val = self._runGather(params, indices)
+    self.assertAllEqual(np.array([[-7, 7], [-7, 7], [-8, 8]]), gather_nd_val)
+
+  def testHigherRankParamsLargerThanIndexSlices(self):
+    params = np.array(
+        [[[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]],
+         [[-80, -10, -20, -30, -70, -50], [80, 10, 20, 30, 70, 50]]],
+        dtype=np.float32).T
+    indices = np.array([[4], [4], [0]], np.int32)
+    gather_nd_val = self._runGather(params, indices)
+    self.assertAllEqual(params[[4, 4, 0]], gather_nd_val)
+
+  def testEmptyIndicesLastRankMeansCopyEntireTensor(self):
+    params = np.array(
+        [[[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]],
+         [[-80, -10, -20, -30, -70, -50], [80, 10, 20, 30, 70, 50]]],
+        dtype=np.float32).T
+    indices = np.array([[], []], dtype=np.int32)  # Size (2, 0)
+    gather_nd_val = self._runGather(params, indices)
+    self.assertAllEqual(
+        np.vstack((params[np.newaxis, :], params[np.newaxis, :])),
+        gather_nd_val)
+
+  def testHigherRankParamsAndIndicesLargerThanIndexSlices(self):
+    params = np.array(
+        [[[-8, -1, -2, -3, -7, -5], [8, 1, 2, 3, 7, 5]],
+         [[-80, -10, -20, -30, -70, -50], [80, 10, 20, 30, 70, 50]]],
+        dtype=np.float32).T
+    indices = np.array([[[3], [2], [1]], [[4], [4], [0]]], np.int32)
+    gather_nd_val = self._runGather(params, indices)
+    self.assertAllEqual(params[[3, 2, 1, 4, 4, 0]].reshape(2, 3, 2, 2),
+                        gather_nd_val)
+
+  def testHigherRankParams(self):
+    shape = (10, 20, 5, 1, 17)
+    params = np.random.rand(*shape).astype(np.float32)
+    indices = np.vstack(
+        [np.random.randint(0, s, size=2000, dtype=np.int32) for s in shape]).T
+    gather_nd_val = self._runGather(params, indices)
+
+    expected = params[tuple(indices.T)]
+    self.assertAllEqual(expected, gather_nd_val)
+
+  def testHigherRankParamsAndIndices(self):
+    shape = (10, 20, 5, 1, 17)
+    params = np.random.rand(*shape).astype(np.float32)
+    indices = np.vstack(
+        [np.random.randint(0, s, size=2000, dtype=np.int32) for s in shape]).T
+    indices_reshaped = indices.reshape([10, 10, 20, 5])
+    gather_nd_val = self._runGather(params, indices_reshaped)
+    expected = params[tuple(indices.T)]
+    self.assertAllEqual(expected.reshape([10, 10, 20]), gather_nd_val)
+
+  def testIgnoreBadIndices(self):
+    shape = (3, 4, 5)
+    params = np.arange(np.prod(shape), dtype=np.int32).reshape(shape)
+    indices = np.array([[[0, 0], [-1, 3]], [[2, 4], [2, 3]]], dtype=np.int32)
+    gather_nd_val = self._runGather(
+        params, indices, bad_indices_policy="IGNORE"
+    )
+    expected = np.array(
+        [
+            [[0, 1, 2, 3, 4], [0, 0, 0, 0, 0]],
+            [[0, 0, 0, 0, 0], [55, 56, 57, 58, 59]],
+        ],
+        dtype=np.int32,
+    )
+    self.assertAllEqual(expected, gather_nd_val)
+
+  def testIgnoreBadIndicesNoFeatureDim(self):
+    shape = (3, 4)
+    params = np.arange(np.prod(shape), dtype=np.int32).reshape(shape)
+    indices = np.array([[[0, 0], [-1, 3]], [[2, 4], [2, 3]]], dtype=np.int32)
+    gather_nd_val = self._runGather(
+        params, indices, bad_indices_policy="IGNORE"
+    )
+    expected = np.array(
+        [[0, 0], [0, 11]],
+        dtype=np.int32,
+    )
+    self.assertAllEqual(expected, gather_nd_val)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/gather_test.py

@@ -0,0 +1,217 @@
+# 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 Gather Op."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import test
+
+FLAGS = flags.FLAGS
+
+
+class GatherTest(xla_test.XLATestCase):
+
+  def _buildParams(self, data, dtype):
+    data = data.astype(dtype.as_numpy_dtype)
+    # For complex types, adds an index-dependent imaginary component so we can
+    # tell we got the right value.
+    if dtype.is_complex:
+      return data + 10j * data
+    return data
+
+  def testScalar1D(self):
+    with self.session() as session, self.test_scope():
+      data = np.array([0, 1, 2, 3, 7, 5])
+      for dtype in self.all_tf_types:
+        for indices in 4, [4], [1, 2, 2, 4, 5]:
+          params_np = self._buildParams(data, dtype)
+          params = array_ops.placeholder(dtype=dtype)
+          indices_tf = constant_op.constant(indices)
+          gather_t = array_ops.gather(params, indices_tf)
+          gather_val = session.run(gather_t, feed_dict={params: params_np})
+          np_val = constant_op.constant(params_np[indices])
+          self.assertAllEqual(np_val, gather_val)
+
+  def testScalar2D(self):
+    with self.session() as session, self.test_scope():
+      data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
+                       [12, 13, 14]])
+      for dtype in self.all_tf_types:
+        for axis in 0, 1, -1:
+          params_np = self._buildParams(data, dtype)
+          params = array_ops.placeholder(dtype=dtype)
+          indices = constant_op.constant(2)
+          gather_t = array_ops.gather(params, indices, axis=axis)
+          gather_val = session.run(gather_t, feed_dict={params: params_np})
+          expected = constant_op.constant(
+              np.take(params_np, 2, axis=axis), dtype)
+          self.assertAllEqual(expected, gather_val)
+
+  def testSimpleTwoD32(self):
+    with self.session() as session, self.test_scope():
+      data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
+                       [12, 13, 14]])
+      for dtype in self.all_tf_types:
+        for axis in 0, 1, -1:
+          params_np = self._buildParams(data, dtype)
+          params = array_ops.placeholder(dtype=dtype)
+          # The indices must be in bounds for any axis.
+          indices = constant_op.constant([0, 1, 0, 2])
+          gather_t = array_ops.gather(params, indices, axis=axis)
+          gather_val = session.run(gather_t, feed_dict={params: params_np})
+          expected = constant_op.constant(
+              np.take(params_np, [0, 1, 0, 2], axis=axis), dtype)
+          self.assertAllEqual(expected, gather_val)
+
+  def testSimpleTwoD32_Int64Indices(self):
+    if np.int64 not in self.int_types:
+      return
+
+    with self.session() as session, self.test_scope():
+      data = np.array([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11],
+                       [12, 13, 14]])
+      # The indices must be in bounds for any axis.
+      indices_np = np.array([0, 1, 0, 2])
+      for dtype in self.all_tf_types:
+        for axis in 0, 1, -1:
+          params_np = self._buildParams(data, dtype)
+          params = array_ops.placeholder(dtype=dtype)
+          indices = array_ops.placeholder(dtype=dtypes.int64)
+          gather_t = array_ops.gather(params, indices, axis=axis)
+          gather_val = session.run(
+              gather_t, feed_dict={
+                  params: params_np,
+                  indices: indices_np
+              })
+          expected = constant_op.constant(
+              np.take(params_np, [0, 1, 0, 2], axis=axis), dtype)
+          self.assertAllEqual(expected, gather_val)
+
+  def testHigherRank(self):
+    """Check that scalar and empty indices shapes work as well."""
+    shape = (2, 1, 3, 2)
+    for indices_shape in (), (0,), (2, 0), (2, 3):
+      for dtype in self.all_tf_types:
+        for axis in 0, 1, 2, 3, -1, -2:
+          params = self._buildParams(np.random.randn(*shape), dtype)
+          indices = np.random.randint(shape[axis], size=indices_shape)
+          with self.session() as sess, self.test_scope():
+            tf_params = array_ops.placeholder(dtype=dtype)
+            tf_indices = constant_op.constant(indices, dtype=dtypes.int32)
+            gather = array_ops.gather(tf_params, tf_indices, axis=axis)
+            gather_value = sess.run(gather, feed_dict={tf_params: params})
+            gather_np = constant_op.constant(
+                np.take(params, indices, axis=axis), dtype)
+            self.assertAllEqual(gather_np, gather_value)
+
+  def testIndicesWithDifferentDimensions(self):
+    with self.session():
+      for dtype in self.numeric_tf_types:
+        params = array_ops.placeholder(dtype=dtype)
+        indices = array_ops.placeholder(dtype=np.int32)
+        with self.test_scope():
+          gather = array_ops.gather(params, indices)
+        self.assertAllEqual(
+            7, gather.eval(feed_dict={params: [4, 7, 2], indices: 1}))
+        self.assertAllEqual(
+            [7], gather.eval(feed_dict={params: [4, 7, 2], indices: [1]}))
+        self.assertAllEqual(
+            [[7]], gather.eval(feed_dict={params: [4, 7, 2], indices: [[1]]}))
+
+  def testGatherPrecision(self):
+    with self.session() as session, self.test_scope():
+      data = np.array([[0, 0, 0, 0], [0, 2 * (1 + np.exp2(-8)), 0, 0],
+                       [0, 0, 0, 0], [0.015789, 0.0985, 0.55789, 0.3842]])
+      indices = np.array([1, 2, 3, 1])
+      dtype = dtypes.float32
+      params_np = self._buildParams(data, dtype)
+      params = array_ops.placeholder(dtype=dtype)
+      indices_tf = constant_op.constant(indices)
+      gather_t = array_ops.gather(params, indices_tf)
+      gather_val = session.run(gather_t, feed_dict={params: params_np})
+      np_val = params_np[indices]
+      self.assertAllEqual(np_val, gather_val)
+
+
+class GatherBenchmark(test.Benchmark):
+  """Microbenchmarks for the gather op."""
+
+  def _benchmarkGather(self, name, axis, gather_indices, use_xla_jit):
+
+    def BuilderFn():
+      inputs = variables.Variable(
+          array_ops.zeros([100, 100, 10, 100, 50], dtype=dtypes.float32),
+          dtype=dtypes.float32,
+          name='input')
+      indices = variables.Variable(
+          gather_indices, dtype=dtypes.int32, name='indices')
+      gather_t = array_ops.gather(inputs, indices, axis=axis)
+      return '%s.axis%d' % (name, axis), [gather_t]
+
+    xla_test.Benchmark(self, BuilderFn, use_xla_jit=use_xla_jit, device='cpu')
+
+  def _benchmarkSliceGather(self, axis, use_xla_jit):
+    """Benchmarks a gather op that's really a dynamic slice."""
+    self._benchmarkGather('slice_gather', axis, [1], use_xla_jit)
+
+  def _benchmarkNontrivialGather(self, axis, use_xla_jit):
+    self._benchmarkGather('nontrivial_gather', axis, [9, 1, 0, 2] * 4,
+                          use_xla_jit)
+
+  def benchmarkSliceGatherAxis0(self):
+    self._benchmarkSliceGather(axis=0, use_xla_jit=False)
+
+  def benchmarkSliceGatherAxis0XLA(self):
+    self._benchmarkSliceGather(axis=0, use_xla_jit=True)
+
+  def benchmarkSliceGatherAxis1(self):
+    self._benchmarkSliceGather(axis=1, use_xla_jit=False)
+
+  def benchmarkSliceGatherAxis1XLA(self):
+    self._benchmarkSliceGather(axis=1, use_xla_jit=True)
+
+  def benchmarkSliceGatherAxis4(self):
+    self._benchmarkSliceGather(axis=4, use_xla_jit=False)
+
+  def benchmarkSliceGatherAxis4XLA(self):
+    self._benchmarkSliceGather(axis=4, use_xla_jit=True)
+
+  def benchmarkNontrivialGatherAxis0(self):
+    self._benchmarkNontrivialGather(axis=0, use_xla_jit=False)
+
+  def benchmarkNontrivialGatherAxis0XLA(self):
+    self._benchmarkNontrivialGather(axis=0, use_xla_jit=True)
+
+  def benchmarkNontrivialGatherAxis1(self):
+    self._benchmarkNontrivialGather(axis=1, use_xla_jit=False)
+
+  def benchmarkNontrivialGatherAxis1XLA(self):
+    self._benchmarkNontrivialGather(axis=1, use_xla_jit=True)
+
+  def benchmarkNontrivialGatherAxis4(self):
+    self._benchmarkNontrivialGather(axis=4, use_xla_jit=False)
+
+  def benchmarkNontrivialGatherAxis4XLA(self):
+    self._benchmarkNontrivialGather(axis=4, use_xla_jit=True)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/tests/giant_const_op_test.py

@@ -0,0 +1,109 @@
+# Copyright 2021 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 giant const op compilation."""
+
+import os
+import numpy as np
+
+from tensorflow.python.distribute import tpu_strategy as tpu_lib
+from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import remote
+from tensorflow.python.eager import test
+from tensorflow.python.framework import config
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.platform import flags
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("tpu", "", "Name of TPU to connect to.")
+flags.DEFINE_string("project", None, "Name of GCP project with TPU.")
+flags.DEFINE_string("zone", None, "Name of GCP zone with TPU.")
+
+
+def get_tpu_cluster_resolver():
+  resolver = tpu_cluster_resolver.TPUClusterResolver(
+      tpu=FLAGS.tpu,
+      zone=FLAGS.zone,
+      project=FLAGS.project,
+  )
+  return resolver
+
+
+def get_tpu_strategy():
+  resolver = get_tpu_cluster_resolver()
+  remote.connect_to_cluster(resolver)
+  tpu_cluster_resolver.initialize_tpu_system(resolver)
+  return tpu_lib.TPUStrategyV2(resolver)
+
+
+# This test doesn't use XLATestCase like the other tests in this directory.
+# The Const op xla op kernel is compilation only and therefore is not executed
+# with XLA in the on demand compilation mode. Also, here we want to feed the
+# full program to XLA to verify handling of programs with giant constant
+# tensors.
+class GiantConstOp(test.TestCase):
+
+  # Verifies that graphs containing giant const tensors that won't fit in memory
+  # are compiled correctly to HLO.
+  def testGiantConst(self):
+    # Disabling Mlir bridge since using the tf2xla implementation of
+    # StridedSliceop which would get executed in this GiantConst test.
+    config.disable_mlir_bridge()
+    strategy = get_tpu_strategy()
+    types = {
+        dtypes.bool,
+        dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64,
+        dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64,
+        dtypes.float16, dtypes.bfloat16,
+        dtypes.float32, dtypes.float64,
+    }
+    for dtype in types:
+      values = [True if dtype is dtypes.bool else 1]
+
+      if dtype is dtypes.bool:
+        values.append(False)
+      elif dtype is not dtypes.float64:
+        # TPUs don't follow IEEE 754 float64 standard for 64 bit floating point
+        # numbers so it could return different output even with just data
+        # transformation ops without any arithmetic operations.
+        values.extend([dtype.min, dtype.max])
+
+      for value in values:
+
+        @def_function.function
+        def train_step():
+
+          # pylint: disable=cell-var-from-loop
+          def computation():
+            const = constant_op.constant(value, dtype=dtype, shape=[1024]*4)
+            return const[:1, :1, :1, :1]
+
+          return strategy.run(computation, args=())
+
+        output = strategy.experimental_local_results(train_step())[0]
+        expected = np.full((1, 1, 1, 1), value)
+        self.assertAllEqual(output, expected)
+
+if __name__ == "__main__":
+  # Make sure TF_XLA_FLAGS is not already set to avoid dropping the existing
+  # value silently.
+  assert "TF_XLA_FLAGS" not in os.environ
+
+  # Disable tfxla constant folding that always creates full Tensors and will
+  # fail for giant tensors.
+  os.environ["TF_XLA_FLAGS"] = "--tf_xla_disable_constant_folding=true"
+
+  test.main()

tensorflow/compiler/tests/image_ops_jit_compile_test.py

@@ -0,0 +1,61 @@
+# Copyright 2022 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 tensorflow.ops.array_ops.repeat."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class ImageOpsTest(xla_test.XLATestCase):
+
+  def testGradImageResize(self):
+    """Tests that the gradient of image.resize is compilable."""
+    with ops.device("device:{}:0".format(self.device)):
+      img_width = 2048
+      var = variables.Variable(array_ops.ones(1, dtype=dtypes.float32))
+
+      def model(x):
+        x = var * x
+        x = image_ops.resize_images(
+            x,
+            size=[img_width, img_width],
+            method=image_ops.ResizeMethod.BILINEAR)
+        return x
+
+      def train(x, y):
+        with backprop.GradientTape() as tape:
+          output = model(x)
+          loss_value = math_ops.reduce_mean((y - output)**2)
+        grads = tape.gradient(loss_value, [var])
+        return grads
+
+      compiled_train = def_function.function(train, jit_compile=True)
+      x = array_ops.zeros((1, img_width // 2, img_width // 2, 1),
+                          dtype=dtypes.float32)
+      y = array_ops.zeros((1, img_width, img_width, 1), dtype=dtypes.float32)
+      self.assertAllClose(train(x, y), compiled_train(x, y))
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  test.main()

tensorflow/compiler/tests/jit_test.py

@@ -0,0 +1,651 @@
+# 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.
+# ==============================================================================
+"""Tests for JIT compilation on the CPU and GPU devices."""
+
+import os
+
+import numpy as np
+
+from tensorflow.compiler.tests import test_utils
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.compiler.xla import jit
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+
+
+jit_scope = jit.experimental_jit_scope
+
+# Disable rewrites to make sure we don't end up having to update this test
+# whenever we implement new ones.
+def NoRewriteSessionConfig():
+  rewriter_config = rewriter_config_pb2.RewriterConfig(
+      disable_model_pruning=True,
+      arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
+      dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
+      function_optimization=rewriter_config_pb2.RewriterConfig.OFF)
+  graph_options = config_pb2.GraphOptions(rewrite_options=rewriter_config)
+  return config_pb2.ConfigProto(graph_options=graph_options)
+
+
+def CompiledKernel(fn, *inputs, **kwargs):
+  """Execute 'fn' as a compiled XLA kernel, with 'inputs'."""
+  name = kwargs.pop("name", None)
+  noinline = kwargs.pop("noinline", None)
+
+  @function.Defun(func_name=name, noinline=noinline, compiled=True)
+  def Compiled(*args):
+    return fn(*args)
+
+  return Compiled(*inputs)
+
+
+def RunMetadataLabels(run_metadata):
+  """Returns all labels in run_metadata."""
+  labels = []
+  for dev_stats in run_metadata.step_stats.dev_stats:
+    for node_stats in dev_stats.node_stats:
+      labels.append(node_stats.timeline_label)
+  return labels
+
+
+def InLabels(labels, substr):
+  """Returns true iff one of the labels contains substr."""
+  return any(substr in x for x in labels)
+
+
+def MetadataHasXlaRunOp(run_metadata):
+  """Returns true if there are XlaRun kernels in run_metadata's timeline."""
+
+  # TODO(phawkins): find a less hacky way to test whether a kernel ran.
+  return InLabels(RunMetadataLabels(run_metadata), "_XlaRun")
+
+
+class JitLaunchTest(test.TestCase):
+
+  # Evaluates 'fn' on 'args' both directly and as a compiled XLA kernel.
+  # Verifies that the outputs match and that XLA was invoked. 'fn' must take
+  # the same number of tensors as arguments that are in 'args', and must return
+  # a tuple of output tensors.
+  #
+  # If 'require_kernel_launch' is True, then we verify that an XlaCompile/XlaRun
+  # node actually ran. However, it is sometimes possible for XlaCompile/XlaRun
+  # ops to be constant-folded away, so the check is optional.
+  def _compare(self,
+               fn,
+               args,
+               require_kernel_launch=True,
+               name=None,
+               noinline=None):
+    with session_lib.Session(config=NoRewriteSessionConfig()) as sess:
+      placeholders = []
+      feeds = {}
+      for arg in args:
+        placeholder = array_ops.placeholder(
+            dtypes.as_dtype(arg.dtype), list(arg.shape))
+        placeholders.append(placeholder)
+        feeds[placeholder] = arg
+
+      compiled_op = CompiledKernel(
+          fn, *placeholders, name=name, noinline=noinline)
+      direct_op = fn(*placeholders)
+
+      run_metadata = config_pb2.RunMetadata()
+      compiled = test_utils.RunWithWarmup(
+          sess, compiled_op, feeds,
+          config_pb2.RunOptions(trace_level=config_pb2.RunOptions.FULL_TRACE),
+          run_metadata)
+      print("Compiled Result {}".format(compiled))
+
+      if require_kernel_launch:
+        self.assertTrue(MetadataHasXlaRunOp(run_metadata))
+
+        direct = sess.run(direct_op, feeds)
+        print("Direct Result {}".format(direct))
+
+        if (isinstance(compiled, (tuple, list)) and
+            (isinstance(direct, (tuple, list)))):
+          for (x, y) in zip(compiled, direct):
+            self.assertAllClose(x, y, rtol=1e-1)
+        else:
+          self.assertAllClose(compiled, direct, rtol=1e-2)
+
+  def testNoOutputs(self):
+    with session_lib.Session() as sess:
+
+      # Check that calling the result as a compiled kernel doesn't crash.
+      @function.Defun(compiled=True)
+      def KernelWithNoOutputs():
+        a = constant_op.constant(100)  # pylint: disable=unused-variable
+
+      call = KernelWithNoOutputs()  # pylint: disable=assignment-from-no-return
+      test_utils.RunWithWarmup(sess, call, {})
+
+  def testAliasing(self):
+    """Regression test for compiled functions that return an aliased buffer.
+
+       XLA returns aliased buffers if outputs are identical. Tests that
+       we handle that case.
+    """
+
+    def AddOnceReturnTwice(x):
+      y = math_ops.add(x, x)
+      return y, y
+
+    # Exercises compiling a function (say, Foo) which calls another function
+    # (say, Bar) which is not inlined. When the compiler compiles Foo, it needs
+    # to symbolically execute Bar correctly regardless of whether Bar is inlined
+    # or not.
+
+    # Tests compiled=True and noinline=True.
+    self._compare(
+        AddOnceReturnTwice, [np.array([[[0.5, -1.0]]], dtype=np.float32)],
+        name="AddOnceReturnTwice_inline",
+        noinline=True)
+
+    # Tests compiled=True and noinline=False.
+    self._compare(
+        AddOnceReturnTwice, [np.array([[[0.5, -1.0]]], dtype=np.float32)],
+        name="AddOnceReturnTwice_noinline",
+        noinline=False)
+
+  def testOneConstOutput(self):
+    """Test consisting of a single constant return value."""
+
+    def OneConstOutput():
+      return constant_op.constant([-3, 44, 99])
+
+    self._compare(OneConstOutput, [], require_kernel_launch=False)
+
+  def testConstZeroElementOutput(self):
+    """Test consisting of a constant zero element return value."""
+
+    def ConstZeroElementOutput():
+      return array_ops.fill([7, 0], 3.0)
+
+    self._compare(ConstZeroElementOutput, [], require_kernel_launch=False)
+
+  def testSomeConstOutputs(self):
+    """Test kernels that return a mixture of const and non-const outputs."""
+
+    def SomeConstOutputs(x):
+      return constant_op.constant(
+          [-2, 7]), array_ops.identity(x), constant_op.constant(3.5)
+
+    self._compare(
+        SomeConstOutputs, [np.array(
+            [[1, 2, 3], [4, 5, 6]], dtype=np.float32)])
+
+  def testInt32Input(self):
+    """Test an int32-typed input.
+
+       On a GPU, int32 tensors will be placed in host memory.
+    """
+
+    def AddToSelf(x):
+      return math_ops.add(x, x)
+
+    self._compare(AddToSelf, [np.array([7, 1, 3], dtype=np.int32)])
+
+  def testMandatoryConstantInput(self):
+    """Tests an operator that has a mandatory-constant shape input."""
+
+    def FillWithFloat(x):
+      return array_ops.fill(x, 9.5)
+
+    self._compare(FillWithFloat, [np.array([3, 2], dtype=np.int32)])
+
+  def testMnistForwardFunc(self):
+    """Compute inference function from MNIST beginners tutorial."""
+    batch_size = 16
+    image_size = 28 * 28
+    num_classes = 10
+
+    # Define a TensorFlow function to compute the forward pass.
+    def MnistForward(w, b, x):
+      return nn_ops.softmax(math_ops.matmul(x, w) + b)
+
+    w = np.random.random_sample((image_size, num_classes)).astype(np.float32)
+    b = np.random.random_sample((num_classes)).astype(np.float32)
+    x = np.random.random_sample((batch_size, image_size)).astype(np.float32)
+    self._compare(MnistForward, [w, b, x])
+
+  def testExplicitMarking(self):
+    """Test explicit marking of operators to compile."""
+    batch_size = 16
+    image_size = 28 * 28
+    num_classes = 10
+
+    with ops.Graph().as_default():
+      x = array_ops.placeholder(dtypes.float32)
+      w = array_ops.placeholder(dtypes.float32)
+      b = array_ops.placeholder(dtypes.float32)
+      with jit_scope():
+        y1 = math_ops.matmul(x, w)
+      y2 = math_ops.add(y1, b)
+      with jit_scope():
+        y = math_ops.square(y2)
+
+      dw = np.random.random_sample((image_size, num_classes)).astype(np.float32)
+      db = np.random.random_sample((num_classes)).astype(np.float32)
+      dx = np.random.random_sample((batch_size, image_size)).astype(np.float32)
+      with session_lib.Session() as sess:
+        run_metadata = config_pb2.RunMetadata()
+        output = test_utils.RunWithWarmup(
+            sess,
+            y, {
+                x: dx,
+                w: dw,
+                b: db
+            },
+            run_metadata=run_metadata,
+            options=config_pb2.RunOptions(
+                trace_level=config_pb2.RunOptions.FULL_TRACE))
+
+        # TODO(phawkins): really we would like to test that there were exactly
+        # two kernel launches. However, we have no reliable way to determine
+        # that.
+        self.assertTrue(MetadataHasXlaRunOp(run_metadata))
+
+        expected = np.square(np.dot(dx, dw) + db)
+        self.assertAllClose(expected, output, rtol=1e-1)
+
+
+class XlaCompilationTest(test.TestCase):
+  """Tests for auto-compilation on CPU/GPU devices."""
+
+  def testReshape(self):
+    """Tests an operator with compile-time constant and non-constant inputs."""
+
+    with self.session(config=NoRewriteSessionConfig()) as sess:
+      x = array_ops.placeholder(dtypes.float32)
+      y = array_ops.placeholder(dtypes.int32)
+      with jit_scope():
+        # Reshape's first argument is non-constant in the JIT, but its second
+        # (shape) argument will be treated as a compile-time constant for
+        # each JIT compilation.
+        # We do not use a tf.const() argument since we want to ensure the
+        # shape is still a run-time argument to the JIT, and not
+        # statically known as part of the JIT compilation's input graph.
+        z = array_ops.reshape(x, y)
+      run_metadata = config_pb2.RunMetadata()
+      out = test_utils.RunWithWarmup(
+          sess,
+          z, {
+              x: np.array([1, 2, 3, 4, 5, 6], np.float32),
+              y: [-1, 3]
+          },
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(MetadataHasXlaRunOp(run_metadata))
+      self.assertAllClose(np.array([[1, 2, 3], [4, 5, 6]], np.float32), out)
+
+  def testIgnoredArguments(self):
+    """Tests that JIT computations can ignore formal parameters."""
+
+    with self.session(config=NoRewriteSessionConfig()) as sess:
+      x = array_ops.placeholder(dtypes.int32)
+      y = array_ops.placeholder(dtypes.int32)
+      with jit_scope():
+        z = math_ops.add(x, x)
+        w = math_ops.add(y, y)
+        # Pulls 'w' into the same compilation via control dependencies.
+        with ops.control_dependencies([w]):
+          n = control_flow_ops.no_op()
+        with ops.control_dependencies([n]):
+          t = math_ops.add(z, z)
+
+      run_metadata = config_pb2.RunMetadata()
+      out = test_utils.RunWithWarmup(
+          sess,
+          t, {
+              x: np.int32(7),
+              y: np.int32(404)
+          },
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(MetadataHasXlaRunOp(run_metadata))
+      self.assertAllClose(28, out)
+
+  def testLoops(self):
+    """Tests that compilation accepts computations containing loops."""
+
+    with self.session(config=NoRewriteSessionConfig()) as session:
+      x = array_ops.placeholder(dtypes.float32)
+      with jit_scope():
+        c = lambda i, _: math_ops.less(i, 5)
+        b = lambda i, x: (i + 1, x * 2.0 + 1.0)
+        _, y = while_loop.while_loop(c, b, (constant_op.constant(0), x))
+
+      run_metadata = config_pb2.RunMetadata()
+      result = session.run(y, {x: np.float32(2)},
+                           run_metadata=run_metadata,
+                           options=config_pb2.RunOptions(
+                               trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(MetadataHasXlaRunOp(run_metadata))
+      self.assertAllClose(result, np.float32(95), rtol=1e-1)
+
+  def testCond(self):
+    """Tests that compilation handles switch operators."""
+
+    with self.session(config=NoRewriteSessionConfig()) as session:
+      x = array_ops.placeholder(dtypes.float32)
+      y = array_ops.placeholder(dtypes.float32)
+      c = array_ops.placeholder(dtypes.bool)
+      with jit_scope():
+        z = x + 1.0
+        w = cond.cond(c, lambda: z, lambda: y)
+        t = math_ops.add(z, w)
+
+      # If JIT compilation chooses to cluster z and t, then execution will
+      # deadlock.
+
+      run_metadata = config_pb2.RunMetadata()
+      result = test_utils.RunWithWarmup(
+          session,
+          t, {
+              x: np.float32(2),
+              y: np.float32(4),
+              c: True
+          },
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(MetadataHasXlaRunOp(run_metadata))
+      self.assertAllClose(result, np.float32(6), rtol=1e-1)
+
+  def testNestedFunction(self):
+    g = ops.Graph()
+    with g.as_default():
+
+      @function.Defun(compiled=True)
+      def Bar(x, y):
+        return x + 2 * y
+
+      @function.Defun(compiled=True)
+      def Foo(x):
+        return Bar(x * x, x * x * x)
+
+      @function.Defun()
+      def Entry(x):
+        return Foo(x)
+
+      inp = array_ops.placeholder(dtypes.float32)
+      out = Entry(inp)
+
+    with self.session(
+        config=NoRewriteSessionConfig(), graph=g, use_gpu=True) as sess:
+      run_metadata = config_pb2.RunMetadata()
+      val = sess.run(out,
+                     feed_dict={inp: [2., 10.]},
+                     run_metadata=run_metadata,
+                     options=config_pb2.RunOptions(
+                         trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertAllClose(val, [20., 2100.])
+
+  def testLoopDeadlock(self):
+    """Regression test for bug that caused deadlocks in graphs with loops."""
+
+    with self.session(config=NoRewriteSessionConfig()) as session:
+      x = array_ops.placeholder(dtypes.float32)
+      with jit_scope():
+        y = x + 1.0
+        c = lambda i, _x, _y: math_ops.less(i, 5)
+        b = lambda i, x, _y: (i + 1, x * 2.0 + 1.0, x - 3.0)
+        _, _, w = while_loop.while_loop(c, b, (constant_op.constant(0), y, x))
+        u = w + y
+      result = session.run(u, {x: np.float32(2)})
+      self.assertAllClose(result, np.float32(63), rtol=1e-1)
+
+  def testGradient(self):
+    """Tests that the backprop function is properly compiled."""
+
+    def _Run(compiled):
+
+      @function.Defun(compiled=compiled)
+      def Forward(x):
+        return math_ops.log(x)
+
+      g = ops.Graph()
+      with g.as_default():
+        x = array_ops.placeholder(dtypes.float32)
+        y = Forward(x)
+        dx, = gradients_impl.gradients(y, [x], 1.0)
+
+      cfg = NoRewriteSessionConfig()
+      cfg.graph_options.optimizer_options.opt_level = (
+          config_pb2.OptimizerOptions.L1)
+      cfg.graph_options.optimizer_options.do_function_inlining = True
+      with session_lib.Session(graph=g, config=cfg) as sess:
+        run_metadata = config_pb2.RunMetadata()
+        dx_val = test_utils.RunWithWarmup(
+            sess,
+            dx,
+            feed_dict={x: 100.},
+            run_metadata=run_metadata,
+            options=config_pb2.RunOptions(
+                trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertAllClose(dx_val, 0.01)
+      return RunMetadataLabels(run_metadata)
+
+    # SymGrad[f=log(x)](x, dy) = 1/x * dy
+    #
+    # Note: we don't need to compute log(x) for dx due to graph pruning.
+
+    # Do not compile the backprop. We should see one Reciprocal and one Mul.
+    labels = _Run(compiled=False)
+    self.assertFalse(InLabels(labels, "Log"))
+    self.assertTrue(InLabels(labels, "Reciprocal"))
+    self.assertTrue(InLabels(labels, "Mul"))
+    self.assertFalse(InLabels(labels, "XlaCompile"))
+    self.assertFalse(InLabels(labels, "XlaRun"))
+
+    # Compile the backprop. One XlaCompile/XlaRun pair.
+    labels = _Run(compiled=True)
+    self.assertFalse(InLabels(labels, "Log"))
+    self.assertFalse(InLabels(labels, "Reciprocal"))
+    self.assertFalse(InLabels(labels, "Mul"))
+    self.assertTrue(InLabels(labels, "XlaCompile"))
+    self.assertTrue(InLabels(labels, "XlaRun"))
+
+
+class ElementWiseFusionTest(test.TestCase):
+
+  # Runs a simple test with the input jit_level and fusion_only flag.
+  def simpleTest(self, arg0, arg1, global_jit_level):
+    config = config_pb2.ConfigProto()
+    config.graph_options.optimizer_options.global_jit_level = global_jit_level
+
+    with session_lib.Session(config=config) as sess:
+      a1 = array_ops.placeholder(dtypes.float32, [2, 2], name="a1")
+      a2 = array_ops.placeholder(dtypes.float32, [2, 2], name="a2")
+      # Two element-wise ops. We need at least two ops since single
+      # element clusters are not passed to XLA in fusion_only mode.
+      a3 = a1 * a2
+      a4 = a3 + a1
+      # A matmul to break XLA clustering.
+      a5 = math_ops.matmul(a4, a1)
+      # Two more element-wise ops.
+      a6 = a5 - a4
+      a7 = a6 + a2
+
+      run_metadata = config_pb2.RunMetadata()
+      output = test_utils.RunWithWarmup(
+          sess,
+          a7, {
+              a1: arg0,
+              a2: arg1
+          },
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+
+      labels = RunMetadataLabels(run_metadata)
+
+      xla_compile_count = sum("XlaCompile(" in x for x in labels)
+      xla_run_count = sum("XlaRun(" in x for x in labels)
+      self.assertEqual(xla_compile_count, xla_run_count)
+
+      return output, xla_run_count
+
+
+class LazyCompilationTest(test.TestCase):
+
+  def testLazyCompilation(self):
+
+    @function.Defun(compiled=True)
+    def CompiledFunction(x):
+      return math_ops.log(x)
+
+    with session_lib.Session(config=NoRewriteSessionConfig()) as sess:
+      x = array_ops.placeholder(dtypes.float32)
+      y = CompiledFunction(x)
+
+      # The very first run of the cluster is always compiled (non-lazily).
+      run_metadata_for_first_run = config_pb2.RunMetadata()
+      sess.run(
+          y,
+          feed_dict={x: [2., 10., 19., 77., 100.]},
+          run_metadata=run_metadata_for_first_run,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(
+          InLabels(
+              RunMetadataLabels(run_metadata_for_first_run), "_XlaCompile"))
+      self.assertTrue(
+          InLabels(RunMetadataLabels(run_metadata_for_first_run), "_XlaRun"))
+
+      run_metadata_before_warmup = config_pb2.RunMetadata()
+      sess.run(
+          y,
+          feed_dict={x: [2., 10.]},
+          run_metadata=run_metadata_before_warmup,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(
+          InLabels(
+              RunMetadataLabels(run_metadata_before_warmup), "_XlaCompile"))
+      self.assertFalse(
+          InLabels(RunMetadataLabels(run_metadata_before_warmup), "_XlaRun"))
+
+      # We compile when we see the same shape a second time.
+
+      run_metadata_after_warmup = config_pb2.RunMetadata()
+      sess.run(
+          y,
+          feed_dict={x: [2., 10.]},
+          run_metadata=run_metadata_after_warmup,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(
+          InLabels(RunMetadataLabels(run_metadata_after_warmup), "_XlaCompile"))
+      self.assertTrue(
+          InLabels(RunMetadataLabels(run_metadata_after_warmup), "_XlaRun"))
+
+      run_metadata_for_new_shape = config_pb2.RunMetadata()
+      sess.run(
+          y,
+          feed_dict={x: [2., 10., 12.]},
+          run_metadata=run_metadata_for_new_shape,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(
+          InLabels(
+              RunMetadataLabels(run_metadata_for_new_shape), "_XlaCompile"))
+      self.assertFalse(
+          InLabels(RunMetadataLabels(run_metadata_for_new_shape), "_XlaRun"))
+
+  def testIsMegamorphic(self):
+
+    @function.Defun(compiled=True)
+    def CompiledFunction(x):
+      return math_ops.log(x)
+
+    with session_lib.Session(config=NoRewriteSessionConfig()) as sess:
+      x = array_ops.placeholder(dtypes.float32)
+      y = CompiledFunction(x)
+
+      # Make the cluster go megamorphic by running it with lots of shape
+      # signatures where the cluster is executed with each signature only a few
+      # times.  Then check that we don't compile the cluster ever again.
+
+      for shape in range(10, 50):
+        for _ in range(0, 49):
+          sess.run(y, feed_dict={x: [0.] * shape})
+
+      for _ in range(0, 50):
+        run_metadata = config_pb2.RunMetadata()
+        sess.run(
+            y,
+            feed_dict={x: [0.] * 60},
+            run_metadata=run_metadata,
+            options=config_pb2.RunOptions(
+                trace_level=config_pb2.RunOptions.FULL_TRACE))
+        self.assertTrue(
+            InLabels(RunMetadataLabels(run_metadata), "_XlaCompile"))
+        self.assertFalse(InLabels(RunMetadataLabels(run_metadata), "_XlaRun"))
+
+  def testIsNotMegamorphic(self):
+
+    @function.Defun(compiled=True)
+    def CompiledFunction(x):
+      return math_ops.log(x)
+
+    with session_lib.Session(config=NoRewriteSessionConfig()) as sess:
+      x = array_ops.placeholder(dtypes.float32)
+      y = CompiledFunction(x)
+
+      # Run the cluster with lots of shape signatures, but in a way that it
+      # isn't megamorphic (i.e. each shape signature sees a lot of executions).
+      # Then check that the cluster has not been marked as megamorphic.
+
+      for shape in range(10, 50):
+        for _ in range(0, 1000):
+          sess.run(y, feed_dict={x: [0.] * shape})
+
+      for _ in range(0, 10):
+        sess.run(y, feed_dict={x: [0.] * 60})
+
+      run_metadata = config_pb2.RunMetadata()
+      sess.run(
+          y,
+          feed_dict={x: [0.] * 60},
+          run_metadata=run_metadata,
+          options=config_pb2.RunOptions(
+              trace_level=config_pb2.RunOptions.FULL_TRACE))
+      self.assertTrue(InLabels(RunMetadataLabels(run_metadata), "_XlaCompile"))
+      self.assertTrue(InLabels(RunMetadataLabels(run_metadata), "_XlaRun"))
+
+
+if __name__ == "__main__":
+  os.environ["TF_XLA_FLAGS"] = ("--tf_xla_enable_lazy_compilation=true " +
+                                os.environ.get("TF_XLA_FLAGS", ""))
+  # This test is using Tensorflow sessions which are not compatible with eager
+  # mode.
+  ops.disable_eager_execution()
+  test.main()

tensorflow/compiler/tests/listdiff_op_test.py

@@ -0,0 +1,96 @@
+# Copyright 2018 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 listdiff operator."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ListDiffTest(xla_test.XLATestCase):
+
+  def _testListDiff(self, x, y, out, idx):
+    for dtype in [dtypes.int32, dtypes.int64]:
+      for index_dtype in [dtypes.int32, dtypes.int64]:
+        with self.session():
+          x_tensor = ops.convert_to_tensor(x, dtype=dtype)
+          y_tensor = ops.convert_to_tensor(y, dtype=dtype)
+          with self.test_scope():
+            out_tensor, idx_tensor = array_ops.listdiff(
+                x_tensor, y_tensor, out_idx=index_dtype)
+            tf_out, tf_idx = self.evaluate([out_tensor, idx_tensor])
+        self.assertAllEqual(out, tf_out)
+        self.assertAllEqual(idx, tf_idx)
+        self.assertEqual(1, out_tensor.get_shape().ndims)
+        self.assertEqual(1, idx_tensor.get_shape().ndims)
+
+  def testBasic1(self):
+    self._testListDiff(x=[1, 2, 3, 4], y=[1, 2], out=[3, 4], idx=[2, 3])
+
+  def testBasic2(self):
+    self._testListDiff(x=[1, 2, 3, 4], y=[2], out=[1, 3, 4], idx=[0, 2, 3])
+
+  def testBasic3(self):
+    self._testListDiff(x=[1, 4, 3, 2], y=[4, 2], out=[1, 3], idx=[0, 2])
+
+  def testDuplicates(self):
+    self._testListDiff(x=[1, 2, 4, 3, 2, 3, 3, 1],
+                       y=[4, 2],
+                       out=[1, 3, 3, 3, 1],
+                       idx=[0, 3, 5, 6, 7])
+
+  def testRandom(self):
+    num_random_tests = 10
+    int_low = -7
+    int_high = 8
+    max_size = 50
+    for _ in range(num_random_tests):
+      x_size = np.random.randint(max_size + 1)
+      x = np.random.randint(int_low, int_high, size=x_size)
+      y_size = np.random.randint(max_size + 1)
+      y = np.random.randint(int_low, int_high, size=y_size)
+      out_idx = [(entry, pos) for pos, entry in enumerate(x) if entry not in y]
+      if out_idx:
+        out, idx = map(list, zip(*out_idx))
+      else:
+        out = []
+        idx = []
+      self._testListDiff(list(x), list(y), out, idx)
+
+  def testFullyOverlapping(self):
+    self._testListDiff(x=[1, 2, 3, 4], y=[1, 2, 3, 4], out=[], idx=[])
+
+  def testNonOverlapping(self):
+    self._testListDiff(x=[1, 2, 3, 4],
+                       y=[5, 6],
+                       out=[1, 2, 3, 4],
+                       idx=[0, 1, 2, 3])
+
+  def testEmptyX(self):
+    self._testListDiff(x=[], y=[1, 2], out=[], idx=[])
+
+  def testEmptyY(self):
+    self._testListDiff(x=[1, 2, 3, 4], y=[], out=[1, 2, 3, 4], idx=[0, 1, 2, 3])
+
+  def testEmptyXY(self):
+    self._testListDiff(x=[], y=[], out=[], idx=[])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/lrn_ops_test.py

@@ -0,0 +1,137 @@
+# 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.
+# ==============================================================================
+"""Tests for Local Response Normalization ops."""
+
+import copy
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import nn
+from tensorflow.python.platform import googletest
+
+CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0"
+
+
+# Local response normalization tests. The forward tests are copied from
+# tensorflow/python/kernel_tests/lrn_op_test.py
+class LRNTest(xla_test.XLATestCase):
+
+  def _LRN(self, input_image, lrn_depth_radius=5, bias=1.0, alpha=1.0,
+           beta=0.5):
+    """Compute expected result."""
+    output = copy.deepcopy(input_image)
+    batch_size = input_image.shape[0]
+    rows = input_image.shape[1]
+    cols = input_image.shape[2]
+    depth = input_image.shape[3]
+    for b in range(batch_size):
+      for r in range(rows):
+        for c in range(cols):
+          for d in range(depth):
+            begin = max(0, d - lrn_depth_radius)
+            end = min(depth, d + lrn_depth_radius + 1)
+            patch = input_image[b, r, c, begin:end]
+            output[b, r, c, d] /= (
+                np.power(bias + alpha * np.sum(patch * patch), beta))
+    return output
+
+  def _RunAndVerify(self, dtype):
+    with self.session():
+      # random shape
+      shape = np.random.randint(1, 16, size=4)
+      # Make depth at least 2 to make it meaningful
+      shape[3] += 1
+      p = array_ops.placeholder(dtype, shape=shape)
+      # random depth_radius, bias, alpha, beta
+      lrn_depth_radius = np.random.randint(1, shape[3])
+      bias = 1.0 + np.random.rand()
+      alpha = 2.0 * np.random.rand()
+      beta = 2.0 * np.random.rand()
+      with self.test_scope():
+        lrn_t = nn.local_response_normalization(
+            p,
+            name="lrn",
+            depth_radius=lrn_depth_radius,
+            bias=bias,
+            alpha=alpha,
+            beta=beta)
+      params = {p: np.random.rand(*shape).astype("f")}
+      result = lrn_t.eval(feed_dict=params)
+    expected = self._LRN(
+        params[p],
+        lrn_depth_radius=lrn_depth_radius,
+        bias=bias,
+        alpha=alpha,
+        beta=beta)
+    err = np.amax(np.abs(result - expected))
+    print("LRN error for bias ", bias, "alpha ", alpha, " beta ", beta, " is ",
+          err)
+    if dtype == dtypes.float32:
+      self.assertTrue(err < 1e-4)
+    else:
+      self.assertTrue(err < 1e-2)
+    self.assertShapeEqual(expected, lrn_t)
+
+  def testCompute(self):
+    for dtype in [dtypes.float32, dtypes.float16, dtypes.bfloat16]:
+      for _ in range(2):
+        self._RunAndVerify(dtype)
+
+  def testLrnGrad(self):
+    # Test for LRNGrad that compares against the CPU implementation.
+    # Note: TF CPU implementation for LRNGrad supports float32 and float16.
+    for dtype in [dtypes.float32, dtypes.float16]:
+      shape = [1, 2, 3, 4]
+      total_size = np.prod(shape)
+      np_dtype = dtype.as_numpy_dtype
+      in_image_vals = np.arange(1, total_size + 1, dtype=np_dtype)
+      out_image_vals = np.arange(1, total_size + 1, dtype=np_dtype)
+      out_grads_vals = np.arange(1, total_size + 1, dtype=np_dtype)
+      depth_radius = np.random.randint(1, shape[3])
+      bias = 1.0 + np.random.rand()
+      alpha = 1.0 * np.random.rand()
+      beta = 1.0 * np.random.rand()
+
+      with self.session():
+        in_image = constant_op.constant(in_image_vals, shape=shape, dtype=dtype)
+        out_image = constant_op.constant(
+            out_image_vals, shape=shape, dtype=dtype
+        )
+        out_grads = constant_op.constant(
+            out_grads_vals, shape=shape, dtype=dtype
+        )
+        with ops.device(CPU_DEVICE):
+          expected = gen_nn_ops.lrn_grad(
+              out_grads, in_image, out_image, depth_radius, bias, alpha, beta
+          )
+        with self.test_scope():
+          actual = gen_nn_ops.lrn_grad(
+              out_grads, in_image, out_image, depth_radius, bias, alpha, beta
+          )
+        expected_val = self.evaluate(expected)
+        actual_val = self.evaluate(actual)
+      rtol = 1e-3 if dtype == dtypes.float32 else 1e-2
+      atol = 1e-6 if dtype == dtypes.float32 else 2e-3
+      self.assertAllClose(actual_val, expected_val, rtol=rtol, atol=atol)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/lstm.py

@@ -0,0 +1,156 @@
+# 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.
+# ==============================================================================
+"""A simple LSTM layer with benchmarks.
+
+This sets up a simple LSTM (Long Short Term Memory) layer, unrolled to a fixed
+length sequence.  The only deviation from standard LSTM cells is that
+activations are clipped, inspired by the GNMT machine translation model.
+The GNMT paper has more details: https://arxiv.org/abs/1609.08144
+"""
+
+from six.moves import range
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_v1
+
+
+def Clip(x):
+  """Clips x to the range [-1., 1.]."""
+  return math_ops.maximum(math_ops.minimum(x, 1.), -1.)
+
+
+def LSTMCellWeightsShape(num_inputs, num_nodes):
+  """Returns the shape of the weights for a single LSTM cell."""
+  # Dimension 0 accounts for combining x with the previous m state.
+  # Dimension 1 accounts for the in value and the (in, forget, out) gates.
+  return [num_inputs + num_nodes, 4 * num_nodes]
+
+
+def LSTMCell(weights, m_prev, c_prev, x, pad):
+  """Unrolls a single LSTM cell with clipped activations forward by one step.
+
+  Args:
+    weights: Weight matrix with shape LSTMCellWeightsShape.
+    m_prev: Previous m states with shape [batch_size, num_nodes].
+    c_prev: Previous c states with shape [batch_size, num_nodes].
+    x: Input with shape [batch_size, num_inputs].
+    pad: Padding with shape [batch_size, 1].  Each padding value is either
+        0 or 1, where 1 indicates padding; i.e. the input is shorter than the
+        sequence length, and the (m, c) states should simply be passed through
+        from the previous states.
+  Returns:
+    The next (m, c) states, each with shape [batch_size, num_nodes].
+  """
+  # Apply weights to the input and previous hidden state.
+  # The matmul here is the "big" operation.
+  xm = array_ops.concat([x, m_prev], 1)
+  xmw = math_ops.matmul(xm, weights)
+
+  # Element-wise ops for the standard LSTM cell, with clipped activations.
+  # XLA can fuse these operations into a single loop.
+  in_value, in_gate, forget_gate, out_gate = array_ops.split(
+      value=xmw, num_or_size_splits=4, axis=1)
+  in_value = math_ops.tanh(in_value)
+  in_gate = math_ops.sigmoid(in_gate)
+  forget_gate = math_ops.sigmoid(forget_gate)
+  out_gate = math_ops.sigmoid(out_gate)
+  c_next = Clip(Clip(forget_gate * c_prev) + Clip(in_gate * in_value))
+  m_next = Clip(out_gate * c_next)
+
+  # Account for padding.
+  c_next = c_prev * pad + c_next * (1.0 - pad)
+  m_next = m_prev * pad + m_next * (1.0 - pad)
+
+  return m_next, c_next
+
+
+def LSTMLayer(cell_name, weights, m, c, x_seq, pad_seq):
+  """Unrolls a layer of LSTM cells forward by the sequence length.
+
+  The sequence length is determined by the length of x_seq and pad_seq, which
+  must be the same.
+
+  Args:
+    cell_name: Base name of each cell.
+    weights: Weight matrix with shape LSTMCellWeightsShape.
+    m: Initial m states with shape [batch_size, num_nodes].
+    c: Initial c states with shape [batch_size, num_nodes].
+    x_seq: List of inputs, each with shape [batch_size, num_inputs].
+        The length of the list is the sequence length.
+    pad_seq: List of paddings, each with shape [batch_size, 1].
+        The length of the list is the sequence length.
+        Each padding value is either 0 or 1, where 1 indicates padding;
+        i.e. the input is shorter than the sequence length.
+  Returns:
+    List of per-sequence-step outputs, each with shape [batch_size, num_nodes].
+  Raises:
+    ValueError: If len(x_seq) != len(pad_seq).
+  """
+  if len(x_seq) != len(pad_seq):
+    raise ValueError('length of x_seq(%d) != pad_seq(%d)' %
+                     (len(x_seq), len(pad_seq)))
+  out_seq = []
+  for seq in range(len(x_seq)):
+    with ops.name_scope('%s_%d' % (cell_name, seq)):
+      m, c = LSTMCell(weights, m, c, x_seq[seq], pad_seq[seq])
+      out_seq.append(array_ops.identity(m, name='out'))
+  return out_seq
+
+
+def RandomVar(shape, name=None):
+  """Returns a variable of the given shape initialized to random values."""
+  return variable_v1.VariableV1(
+      random_ops.random_uniform(shape), dtype=dtypes.float32, name=name)
+
+
+def RandomInputs(batch_size, seq_length, num_inputs):
+  """Returns randomly initialized (x_seq, pad_seq) sequences."""
+  x_seq = []
+  pad_seq = []
+  with ops.name_scope('inputs'):
+    for seq in range(seq_length):
+      x_seq.append(RandomVar([batch_size, num_inputs], name='x_seq_%d' % seq))
+      # Real padding values are always a sequence of 0 followed by a
+      # sequence of 1, but random values are fine for benchmarking.
+      pad_seq.append(RandomVar([batch_size, 1], name='pad_seq_%d' % seq))
+  return x_seq, pad_seq
+
+
+def BuildLSTMLayer(batch_size, seq_length, num_inputs, num_nodes):
+  """Builds a single LSTM layer with random weights and inputs.
+
+  Args:
+    batch_size: Inputs are fed in batches of this size.
+    seq_length: The sequence length to unroll the LSTM layer.
+    num_inputs: Dimension of inputs that are fed into each LSTM cell.
+    num_nodes: The number of nodes in each LSTM cell.
+
+  Returns:
+    (out_seq, weights) pair.  The out_seq is a list of per-sequence-step
+    outputs, each with shape [batch_size, num_nodes].  The weights are a list of
+    weight variables that may be trained.
+  """
+  weights = RandomVar(
+      LSTMCellWeightsShape(num_inputs, num_nodes), name='weights')
+  m = array_ops.zeros([batch_size, num_nodes], name='init_m')
+  c = array_ops.zeros([batch_size, num_nodes], name='init_c')
+  x_seq, pad_seq = RandomInputs(batch_size, seq_length, num_inputs)
+
+  out_seq = LSTMLayer('lstm', weights, m, c, x_seq, pad_seq)
+  return out_seq, [weights]

tensorflow/compiler/tests/lstm_layer_inference.config.pbtxt

@@ -0,0 +1,20 @@
+# Text form of tensorflow.tf2xla.Config proto.
+feed{ id{node_name:"inputs/x_seq_0/read"} shape{dim{size:128}dim{size:1024}} }
+feed{ id{node_name:"inputs/x_seq_1/read"} shape{dim{size:128}dim{size:1024}} }
+feed{ id{node_name:"inputs/x_seq_2/read"} shape{dim{size:128}dim{size:1024}} }
+feed{ id{node_name:"inputs/x_seq_3/read"} shape{dim{size:128}dim{size:1024}} }
+feed{ id{node_name:"inputs/x_seq_4/read"} shape{dim{size:128}dim{size:1024}} }
+feed{ id{node_name:"inputs/pad_seq_0/read"} shape{dim{size:128}dim{size:1}} }
+feed{ id{node_name:"inputs/pad_seq_1/read"} shape{dim{size:128}dim{size:1}} }
+feed{ id{node_name:"inputs/pad_seq_2/read"} shape{dim{size:128}dim{size:1}} }
+feed{ id{node_name:"inputs/pad_seq_3/read"} shape{dim{size:128}dim{size:1}} }
+feed{ id{node_name:"inputs/pad_seq_4/read"} shape{dim{size:128}dim{size:1}} }
+feed{ id{node_name:"weights/read"} shape{dim{size:2048}dim{size:4096}} }
+feed{ id{node_name:"init_c"} shape{dim{size:128}dim{size:1024}} }
+feed{ id{node_name:"init_m"} shape{dim{size:128}dim{size:1024}} }
+
+fetch{ id{node_name:"lstm_0/out"} }
+fetch{ id{node_name:"lstm_1/out"} }
+fetch{ id{node_name:"lstm_2/out"} }
+fetch{ id{node_name:"lstm_3/out"} }
+fetch{ id{node_name:"lstm_4/out"} }

tensorflow/compiler/tests/lstm_test.py

@@ -0,0 +1,328 @@
+# 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.
+# ==============================================================================
+"""Tests for the LSTM cell and layer."""
+
+import argparse
+import os
+import sys
+
+import numpy as np
+
+from tensorflow.compiler.tests import lstm
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def _DumpGraph(graph, basename):
+  if FLAGS.dump_graph_dir:
+    name = os.path.join(FLAGS.dump_graph_dir, basename + '.pbtxt')
+    with open(name, 'w') as f:
+      f.write(str(graph.as_graph_def()))
+
+
+def _Sigmoid(x):
+  return 1. / (1. + np.exp(-x))
+
+
+def _Clip(x):
+  return np.maximum(np.minimum(x, 1.), -1.)
+
+
+class LSTMTest(test.TestCase):
+
+  def setUp(self):
+    # The tests for a single LSTM cell and LSTM layer use these values as
+    # inputs.  We always set the dimensionality of num_inputs=1; thus batch_size
+    # actually represents the different input cases.
+    self._inputs = np.array([[-1.], [-.5], [0.], [.5], [1.]], np.float32)
+    self._batch_size = len(self._inputs)
+
+  def _NextC(self, inputs, weight, m_prev, c_prev):
+    """Returns the next c states of an LSTM cell."""
+    x = (inputs + m_prev) * weight
+    return _Clip(_Clip(_Sigmoid(x) * c_prev) + _Clip(_Sigmoid(x) * np.tanh(x)))
+
+  def _NextM(self, inputs, weight, m_prev, c_prev):
+    """Returns the next m states of an LSTM cell."""
+    x = (inputs + m_prev) * weight
+    return _Clip(_Sigmoid(x) * self._NextC(inputs, weight, m_prev, c_prev))
+
+  def _RunLSTMCell(self, basename, init_weights, m_prev_scalar, c_prev_scalar,
+                   pad_scalar):
+    with self.session() as sess:
+      num_inputs = 1
+      num_nodes = 1
+
+      weights = init_weights(lstm.LSTMCellWeightsShape(num_inputs, num_nodes))
+      m_prev = constant_op.constant([[m_prev_scalar]] * self._batch_size)
+      c_prev = constant_op.constant([[c_prev_scalar]] * self._batch_size)
+      x = constant_op.constant(self._inputs)
+      pad = constant_op.constant([[pad_scalar]] * self._batch_size)
+
+      m, c = lstm.LSTMCell(weights, m_prev, c_prev, x, pad)
+      _DumpGraph(sess.graph, 'lstm_cell_%s_%d_%d_%d' %
+                 (basename, m_prev_scalar, c_prev_scalar, pad_scalar))
+
+      # Initialize variables and run the unrolled LSTM step.
+      self.evaluate(variables.global_variables_initializer())
+      return self.evaluate([m, c])
+
+  @test_util.run_without_tensor_float_32('TF32 capable devices fail the test'
+                                         ' due to reduced matmul precision')
+  def testLSTMCell(self):
+    # Run with all-0 weights, no padding.
+    m, c = self._RunLSTMCell('zeros', init_ops.zeros_initializer(), 0., 0., 0.)
+    self.assertAllClose(m, [[0.]] * self._batch_size)
+    self.assertAllClose(c, [[0.]] * self._batch_size)
+    m, c = self._RunLSTMCell('zeros', init_ops.zeros_initializer(), 0., 1., 0.)
+    self.assertAllClose(m, [[.25]] * self._batch_size)
+    self.assertAllClose(c, [[.5]] * self._batch_size)
+    m, c = self._RunLSTMCell('zeros', init_ops.zeros_initializer(), 1., 0., 0.)
+    self.assertAllClose(m, [[.0]] * self._batch_size)
+    self.assertAllClose(c, [[.0]] * self._batch_size)
+    m, c = self._RunLSTMCell('zeros', init_ops.zeros_initializer(), 1., 1., 0.)
+    self.assertAllClose(m, [[.25]] * self._batch_size)
+    self.assertAllClose(c, [[.5]] * self._batch_size)
+
+    # Run with all-1 weights, no padding.
+    for m_prev in [0., 1.]:
+      for c_prev in [0., 1.]:
+        m, c = self._RunLSTMCell('ones',
+                                 init_ops.ones_initializer(), m_prev, c_prev,
+                                 0.)
+        self.assertAllClose(m, self._NextM(self._inputs, 1., m_prev, c_prev))
+        self.assertAllClose(c, self._NextC(self._inputs, 1., m_prev, c_prev))
+
+    # Run with random weights.
+    for weight in np.random.rand(3):
+      weight_tf = constant_op.constant(weight, dtypes.float32)
+      random_weight = lambda shape, w=weight_tf: array_ops.fill(shape, w)
+
+      # No padding.
+      for m_prev in [0., 1.]:
+        for c_prev in [0., 1.]:
+          m, c = self._RunLSTMCell('random', random_weight, m_prev, c_prev, 0.)
+          self.assertAllClose(m,
+                              self._NextM(self._inputs, weight, m_prev, c_prev))
+          self.assertAllClose(c,
+                              self._NextC(self._inputs, weight, m_prev, c_prev))
+
+      # Set padding.
+      for m_prev in [0., 1.]:
+        for c_prev in [0., 1.]:
+          m, c = self._RunLSTMCell('random', random_weight, m_prev, c_prev, 1.)
+          self.assertAllClose(m, [[m_prev]] * self._batch_size)
+          self.assertAllClose(c, [[c_prev]] * self._batch_size)
+
+  def testLSTMLayerErrors(self):
+    num_inputs = 1
+    num_nodes = 1
+    seq_length = 3
+
+    weights = array_ops.zeros(lstm.LSTMCellWeightsShape(num_inputs, num_nodes))
+    m = constant_op.constant([[0.]] * self._batch_size)
+    c = constant_op.constant([[0.]] * self._batch_size)
+    x_seq = [constant_op.constant(self._inputs)] * seq_length
+    pad = constant_op.constant([[0.]] * self._batch_size)
+
+    with self.assertRaisesWithPredicateMatch(ValueError, 'length of x_seq'):
+      lstm.LSTMLayer('lstm', weights, m, c, x_seq, [pad])
+    with self.assertRaisesWithPredicateMatch(ValueError, 'length of x_seq'):
+      lstm.LSTMLayer('lstm', weights, m, c, x_seq, [pad] * 2)
+    with self.assertRaisesWithPredicateMatch(ValueError, 'length of x_seq'):
+      lstm.LSTMLayer('lstm', weights, m, c, x_seq, [pad] * 4)
+
+  def _RunLSTMLayer(self, basename, init_weights, m_init_scalar, c_init_scalar,
+                    pad_scalar):
+    with self.session() as sess:
+      num_inputs = 1
+      num_nodes = 1
+      seq_length = 3
+
+      weights = init_weights(lstm.LSTMCellWeightsShape(num_inputs, num_nodes))
+      m_init = constant_op.constant([[m_init_scalar]] * self._batch_size)
+      c_init = constant_op.constant([[c_init_scalar]] * self._batch_size)
+      x_seq = [constant_op.constant(self._inputs)] * seq_length
+      pad_seq = [constant_op.constant([[pad_scalar]] * self._batch_size)
+                ] * seq_length
+
+      out_seq = lstm.LSTMLayer('lstm', weights, m_init, c_init, x_seq, pad_seq)
+      _DumpGraph(sess.graph, 'lstm_layer_%s_%d_%d_%d' %
+                 (basename, m_init_scalar, c_init_scalar, pad_scalar))
+
+      # Initialize variables and run the unrolled LSTM layer.
+      self.evaluate(variables.global_variables_initializer())
+      return self.evaluate(out_seq)
+
+  @test_util.run_without_tensor_float_32('TF32 capable devices fail the test'
+                                         ' due to reduced matmul precision')
+  def testLSTMLayer(self):
+    # Run with all-0 weights, no padding.
+    o = self._RunLSTMLayer('zeros', init_ops.zeros_initializer(), 0., 0., 0.)
+    self.assertAllClose(o, [[[0.]] * self._batch_size] * 3)
+    o = self._RunLSTMLayer('zeros', init_ops.zeros_initializer(), 0., 1., 0.)
+    self.assertAllClose(o, [[[.25]] * self._batch_size,
+                            [[.125]] * self._batch_size,
+                            [[.0625]] * self._batch_size])
+    o = self._RunLSTMLayer('zeros', init_ops.zeros_initializer(), 1., 0., 0.)
+    self.assertAllClose(o, [[[0.]] * self._batch_size] * 3)
+    o = self._RunLSTMLayer('zeros', init_ops.zeros_initializer(), 1., 1., 0.)
+    self.assertAllClose(o, [[[.25]] * self._batch_size,
+                            [[.125]] * self._batch_size,
+                            [[.0625]] * self._batch_size])
+
+    # Run with all-1 weights, no padding.
+    weight1 = 1.
+    for m_init in [0., 1.]:
+      for c_init in [0., 1.]:
+        o = self._RunLSTMLayer('ones',
+                               init_ops.ones_initializer(), m_init, c_init, 0.)
+        m0 = self._NextM(self._inputs, weight1, m_init, c_init)
+        c0 = self._NextC(self._inputs, weight1, m_init, c_init)
+        self.assertAllClose(o[0], m0)
+        m1 = self._NextM(self._inputs, weight1, m0, c0)
+        c1 = self._NextC(self._inputs, weight1, m0, c0)
+        self.assertAllClose(o[1], m1)
+        m2 = self._NextM(self._inputs, weight1, m1, c1)
+        self.assertAllClose(o[2], m2)
+
+    # Run with random weights.
+    for weight in np.random.rand(3):
+      weight_tf = constant_op.constant(weight, dtypes.float32)
+      random_weight = lambda shape, w=weight_tf: array_ops.fill(shape, w)
+
+      # No padding.
+      for m_init in [0., 1.]:
+        for c_init in [0., 1.]:
+          o = self._RunLSTMLayer('random', random_weight, m_init, c_init, 0.)
+          m0 = self._NextM(self._inputs, weight, m_init, c_init)
+          c0 = self._NextC(self._inputs, weight, m_init, c_init)
+          self.assertAllClose(o[0], m0)
+          m1 = self._NextM(self._inputs, weight, m0, c0)
+          c1 = self._NextC(self._inputs, weight, m0, c0)
+          self.assertAllClose(o[1], m1)
+          m2 = self._NextM(self._inputs, weight, m1, c1)
+          self.assertAllClose(o[2], m2)
+
+      # Set padding.
+      o = self._RunLSTMLayer('random', random_weight, 0., 0., 1.)
+      self.assertAllClose(o, [[[0.]] * self._batch_size] * 3)
+      o = self._RunLSTMLayer('random', random_weight, 0., 1., 1.)
+      self.assertAllClose(o, [[[0.]] * self._batch_size] * 3)
+      o = self._RunLSTMLayer('random', random_weight, 1., 0., 1.)
+      self.assertAllClose(o, [[[1.]] * self._batch_size] * 3)
+      o = self._RunLSTMLayer('random', random_weight, 1., 1., 1.)
+      self.assertAllClose(o, [[[1.]] * self._batch_size] * 3)
+
+
+class LSTMBenchmark(test.Benchmark):
+  """Mcro-benchmarks for a single layer of LSTM cells."""
+
+  def _LayerBuilder(self, do_training):
+    out_seq, weights = lstm.BuildLSTMLayer(FLAGS.batch_size, FLAGS.seq_length,
+                                           FLAGS.num_inputs, FLAGS.num_nodes)
+    name, fetches = ('lstm_layer_inference', out_seq)
+    if do_training:
+      # Not a real loss function, but good enough for benchmarking backprop.
+      loss = math_ops.reduce_sum(math_ops.add_n(out_seq))
+      dw = gradients_impl.gradients(loss, weights)
+      name, fetches = ('lstm_layer_training', dw)
+
+    _DumpGraph(ops.get_default_graph(),
+               '%s_%d_%d_%d_%d' % (name, FLAGS.batch_size, FLAGS.seq_length,
+                                   FLAGS.num_inputs, FLAGS.num_nodes))
+    return name, fetches
+
+  def benchmarkLayerInference(self):
+    xla_test.Benchmark(self, lambda: self._LayerBuilder(False), False,
+                       FLAGS.device)
+
+  def benchmarkLayerInferenceXLA(self):
+    xla_test.Benchmark(self, lambda: self._LayerBuilder(False), True,
+                       FLAGS.device)
+
+  def benchmarkLayerTraining(self):
+    xla_test.Benchmark(self, lambda: self._LayerBuilder(True), False,
+                       FLAGS.device)
+
+  def benchmarkLayerTrainingXLA(self):
+    xla_test.Benchmark(self, lambda: self._LayerBuilder(True), True,
+                       FLAGS.device)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.register('type', 'bool', lambda v: v.lower() == 'true')
+  parser.add_argument(
+      '--batch_size',
+      type=int,
+      default=128,
+      help="""\
+      Inputs are fed in batches of this size, for both inference and training.
+      Larger values cause the matmul in each LSTM cell to have higher
+      dimensionality.\
+      """
+  )
+  parser.add_argument(
+      '--seq_length',
+      type=int,
+      default=60,
+      help="""\
+      Length of the unrolled sequence of LSTM cells in a layer.Larger values
+      cause more LSTM matmuls to be run.\
+      """
+  )
+  parser.add_argument(
+      '--num_inputs',
+      type=int,
+      default=1024,
+      help='Dimension of inputs that are fed into each LSTM cell.'
+  )
+  parser.add_argument(
+      '--num_nodes',
+      type=int,
+      default=1024,
+      help='Number of nodes in each LSTM cell.'
+  )
+  parser.add_argument(
+      '--device',
+      type=str,
+      default='gpu',
+      help="""\
+      TensorFlow device to assign ops to, e.g. "gpu", "cpu". For details see
+      documentation for tf.Graph.device.\
+      """
+  )
+  parser.add_argument(
+      '--dump_graph_dir',
+      type=str,
+      default='',
+      help='If non-empty, dump graphs in *.pbtxt format to this directory.'
+  )
+  global FLAGS  # pylint:disable=global-at-module-level
+  FLAGS, unparsed = parser.parse_known_args()
+  # This test is using Tensorflow sessions which are not compatible with eager
+  # mode.
+  ops.disable_eager_execution()
+  test.main(argv=[sys.argv[0]] + unparsed)

tensorflow/compiler/tests/manip_ops_test.py

@@ -0,0 +1,64 @@
+# 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.
+# ==============================================================================
+"""Test cases for manip ops."""
+
+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 manip_ops
+from tensorflow.python.platform import googletest
+
+
+class ManipOpsTest(xla_test.XLATestCase):
+  """Test cases for manip ops."""
+
+  def _testRoll(self, a, shift, axis):
+    with self.session() as session:
+      with self.test_scope():
+        p = array_ops.placeholder(dtypes.as_dtype(a.dtype), a.shape, name="a")
+        output = manip_ops.roll(a, shift, axis)
+      result = session.run(output, {p: a})
+      self.assertAllEqual(result, np.roll(a, shift, axis))
+
+  def testNumericTypes(self):
+    for t in self.numeric_types:
+      self._testRoll(np.random.randint(-100, 100, (5)).astype(t), 3, 0)
+      self._testRoll(
+          np.random.randint(-100, 100, (4, 4, 3)).astype(t), [1, -6, 6],
+          [0, 1, 2])
+      self._testRoll(
+          np.random.randint(-100, 100, (4, 2, 1, 3)).astype(t), [0, 1, -2],
+          [1, 2, 3])
+
+  def testFloatTypes(self):
+    for t in self.float_types:
+      self._testRoll(np.random.rand(5).astype(t), 2, 0)
+      self._testRoll(np.random.rand(3, 4).astype(t), [1, 2], [1, 0])
+      self._testRoll(np.random.rand(1, 3, 4).astype(t), [1, 0, -3], [0, 1, 2])
+
+  def testComplexTypes(self):
+    for t in self.complex_types:
+      x = np.random.rand(4, 4).astype(t)
+      self._testRoll(x + 1j * x, 2, 0)
+      x = np.random.rand(2, 5).astype(t)
+      self._testRoll(x + 1j * x, [1, 2], [1, 0])
+      x = np.random.rand(3, 2, 1, 1).astype(t)
+      self._testRoll(x + 1j * x, [2, 1, 1, 0], [0, 3, 1, 2])
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/matrix_band_part_test.py

@@ -0,0 +1,195 @@
+# Copyright 2018 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.
+# ==============================================================================
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+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 test
+
+
+class MatrixBandPartTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      {
+          'batch_shape': [],
+          'rows': 1,
+          'cols': 1
+      },
+      {
+          'batch_shape': [],
+          'rows': 1,
+          'cols': 2
+      },
+      {
+          'batch_shape': [],
+          'rows': 1,
+          'cols': 7
+      },
+      {
+          'batch_shape': [],
+          'rows': 2,
+          'cols': 1
+      },
+      {
+          'batch_shape': [],
+          'rows': 2,
+          'cols': 2
+      },
+      {
+          'batch_shape': [],
+          'rows': 2,
+          'cols': 7
+      },
+      {
+          'batch_shape': [],
+          'rows': 7,
+          'cols': 1
+      },
+      {
+          'batch_shape': [],
+          'rows': 7,
+          'cols': 2
+      },
+      {
+          'batch_shape': [],
+          'rows': 7,
+          'cols': 7
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 1,
+          'cols': 1
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 1,
+          'cols': 2
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 1,
+          'cols': 7
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 2,
+          'cols': 1
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 2,
+          'cols': 2
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 2,
+          'cols': 7
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 7,
+          'cols': 1
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 7,
+          'cols': 2
+      },
+      {
+          'batch_shape': [2,],
+          'rows': 7,
+          'cols': 7
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 1,
+          'cols': 1
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 1,
+          'cols': 2
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 1,
+          'cols': 7
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 2,
+          'cols': 1
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 2,
+          'cols': 2
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 2,
+          'cols': 7
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 7,
+          'cols': 1
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 7,
+          'cols': 2
+      },
+      {
+          'batch_shape': [1, 3, 2],
+          'rows': 7,
+          'cols': 7
+      },
+  )
+  def testMatrixBandPart(self, batch_shape, rows, cols):
+    # TODO(b/125505881): Disabled due to LLVM backend crash.
+    if self.device == 'XLA_CPU' and cols == 7 and rows == 1 and batch_shape == [
+        1, 3, 2
+    ]:
+      pass
+    for dtype in self.float_types:
+      with self.session():
+        mat = np.ones(batch_shape + [rows, cols]).astype(dtype)
+        batch_mat = np.tile(mat, batch_shape + [1, 1])
+        for lower in -1, 0, 1, rows - 1:
+          for upper in -1, 0, 1, cols - 1:
+            band_np = mat
+            if lower >= 0:
+              band_np = np.triu(band_np, -lower)
+            if upper >= 0:
+              band_np = np.tril(band_np, upper)
+            if batch_shape:
+              band_np = np.tile(band_np, batch_shape + [1, 1])
+
+            placeholder = array_ops.placeholder(dtype)
+            with self.test_scope():
+              band = array_ops.matrix_band_part(
+                  placeholder, constant_op.constant(lower, dtype=dtypes.int32),
+                  constant_op.constant(upper, dtype=dtypes.int32))
+              feed_dict = {placeholder: batch_mat}
+              self.assertAllEqual(band_np, band.eval(feed_dict=feed_dict))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/matrix_diag_ops_test.py

@@ -0,0 +1,772 @@
+# 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 XLA matrix diag ops."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.platform import googletest
+
+default_v2_alignment = "LEFT_LEFT"
+alignment_list = ["RIGHT_LEFT", "LEFT_RIGHT"]
+
+
+def zip_to_first_list_length(a, b):
+  if len(b) > len(a):
+    return zip(a, b[:len(a)])
+  return zip(a, b + [None] * (len(a) - len(b)))
+
+
+# Routines to convert test cases to have diagonals in a specified alignment.
+# Copied from //third_party/tensorflow/python/kernel_tests/array_ops/
+# diag_op_test.py
+def repack_diagonals(packed_diagonals,
+                     diag_index,
+                     num_rows,
+                     num_cols,
+                     align=None):
+  # The original test cases are LEFT_LEFT aligned.
+  if align == default_v2_alignment or align is None:
+    return packed_diagonals
+
+  align = align.split("_")
+  d_lower, d_upper = diag_index
+  batch_dims = packed_diagonals.ndim - (2 if d_lower < d_upper else 1)
+  max_diag_len = packed_diagonals.shape[-1]
+  index = (slice(None),) * batch_dims
+  repacked_diagonals = np.zeros_like(packed_diagonals)
+
+  # Aligns each diagonal row-by-row.
+  for diag_index in range(d_lower, d_upper + 1):
+    diag_len = min(num_rows + min(0, diag_index), num_cols - max(0, diag_index))
+    row_index = d_upper - diag_index
+    padding_len = max_diag_len - diag_len
+    left_align = (diag_index >= 0 and
+                  align[0] == "LEFT") or (diag_index <= 0 and
+                                          align[1] == "LEFT")
+    # Prepares index tuples.
+    extra_dim = tuple() if d_lower == d_upper else (row_index,)
+    packed_last_dim = (slice(None),) if left_align else (slice(0, diag_len, 1),)
+    repacked_last_dim = (slice(None),) if left_align else (slice(
+        padding_len, max_diag_len, 1),)
+    packed_index = index + extra_dim + packed_last_dim
+    repacked_index = index + extra_dim + repacked_last_dim
+
+    # Repacks the diagonal.
+    repacked_diagonals[repacked_index] = packed_diagonals[packed_index]
+  return repacked_diagonals
+
+
+def repack_diagonals_in_tests(tests, align=None):
+  # The original test cases are LEFT_LEFT aligned.
+  if align == default_v2_alignment or align is None:
+    return tests
+
+  new_tests = dict()
+  # Loops through each case.
+  for diag_index, (packed_diagonals, padded_diagonals) in tests.items():
+    num_rows, num_cols = padded_diagonals.shape[-2:]
+    repacked_diagonals = repack_diagonals(
+        packed_diagonals, diag_index, num_rows, num_cols, align=align)
+    new_tests[diag_index] = (repacked_diagonals, padded_diagonals)
+
+  return new_tests
+
+
+# Test cases shared by MatrixDiagV2, MatrixDiagPartV2, and MatrixSetDiagV2.
+# Copied from //third_party/tensorflow/python/kernel_tests/array_ops/
+# diag_op_test.py
+def square_cases(align=None):
+  # pyformat: disable
+  mat = np.array([[[1, 2, 3, 4, 5],
+                   [6, 7, 8, 9, 1],
+                   [3, 4, 5, 6, 7],
+                   [8, 9, 1, 2, 3],
+                   [4, 5, 6, 7, 8]],
+                  [[9, 1, 2, 3, 4],
+                   [5, 6, 7, 8, 9],
+                   [1, 2, 3, 4, 5],
+                   [6, 7, 8, 9, 1],
+                   [2, 3, 4, 5, 6]]])
+  tests = dict()
+  # tests[d_lower, d_upper] = (compact_diagonals, padded_diagonals)
+  tests[-1, -1] = (np.array([[6, 4, 1, 7],
+                             [5, 2, 8, 5]]),
+                   np.array([[[0, 0, 0, 0, 0],
+                              [6, 0, 0, 0, 0],
+                              [0, 4, 0, 0, 0],
+                              [0, 0, 1, 0, 0],
+                              [0, 0, 0, 7, 0]],
+                             [[0, 0, 0, 0, 0],
+                              [5, 0, 0, 0, 0],
+                              [0, 2, 0, 0, 0],
+                              [0, 0, 8, 0, 0],
+                              [0, 0, 0, 5, 0]]]))
+  tests[-4, -3] = (np.array([[[8, 5],
+                              [4, 0]],
+                             [[6, 3],
+                              [2, 0]]]),
+                   np.array([[[0, 0, 0, 0, 0],
+                              [0, 0, 0, 0, 0],
+                              [0, 0, 0, 0, 0],
+                              [8, 0, 0, 0, 0],
+                              [4, 5, 0, 0, 0]],
+                             [[0, 0, 0, 0, 0],
+                              [0, 0, 0, 0, 0],
+                              [0, 0, 0, 0, 0],
+                              [6, 0, 0, 0, 0],
+                              [2, 3, 0, 0, 0]]]))
+  tests[-2, 1] = (np.array([[[2, 8, 6, 3, 0],
+                             [1, 7, 5, 2, 8],
+                             [6, 4, 1, 7, 0],
+                             [3, 9, 6, 0, 0]],
+                            [[1, 7, 4, 1, 0],
+                             [9, 6, 3, 9, 6],
+                             [5, 2, 8, 5, 0],
+                             [1, 7, 4, 0, 0]]]),
+                  np.array([[[1, 2, 0, 0, 0],
+                             [6, 7, 8, 0, 0],
+                             [3, 4, 5, 6, 0],
+                             [0, 9, 1, 2, 3],
+                             [0, 0, 6, 7, 8]],
+                            [[9, 1, 0, 0, 0],
+                             [5, 6, 7, 0, 0],
+                             [1, 2, 3, 4, 0],
+                             [0, 7, 8, 9, 1],
+                             [0, 0, 4, 5, 6]]]))
+  tests[2, 4] = (np.array([[[5, 0, 0],
+                            [4, 1, 0],
+                            [3, 9, 7]],
+                           [[4, 0, 0],
+                            [3, 9, 0],
+                            [2, 8, 5]]]),
+                 np.array([[[0, 0, 3, 4, 5],
+                            [0, 0, 0, 9, 1],
+                            [0, 0, 0, 0, 7],
+                            [0, 0, 0, 0, 0],
+                            [0, 0, 0, 0, 0]],
+                           [[0, 0, 2, 3, 4],
+                            [0, 0, 0, 8, 9],
+                            [0, 0, 0, 0, 5],
+                            [0, 0, 0, 0, 0],
+                            [0, 0, 0, 0, 0]]]))
+  # pyformat: enable
+  return (mat, repack_diagonals_in_tests(tests, align))
+
+
+def tall_cases(align=None):
+  # pyformat: disable
+  mat = np.array([[[1, 2, 3],
+                   [4, 5, 6],
+                   [7, 8, 9],
+                   [9, 8, 7],
+                   [6, 5, 4]],
+                  [[3, 2, 1],
+                   [1, 2, 3],
+                   [4, 5, 6],
+                   [7, 8, 9],
+                   [9, 8, 7]]])
+  tests = dict()
+  # tests[d_lower, d_upper] = (compact_diagonals, padded_diagonals)
+  tests[0, 0] = (np.array([[1, 5, 9],
+                           [3, 2, 6]]),
+                 np.array([[[1, 0, 0],
+                            [0, 5, 0],
+                            [0, 0, 9],
+                            [0, 0, 0]],
+                           [[3, 0, 0],
+                            [0, 2, 0],
+                            [0, 0, 6],
+                            [0, 0, 0]]]))
+  tests[-4, -3] = (np.array([[[9, 5],
+                              [6, 0]],
+                             [[7, 8],
+                              [9, 0]]]),
+                   np.array([[[0, 0, 0],
+                              [0, 0, 0],
+                              [0, 0, 0],
+                              [9, 0, 0],
+                              [6, 5, 0]],
+                             [[0, 0, 0],
+                              [0, 0, 0],
+                              [0, 0, 0],
+                              [7, 0, 0],
+                              [9, 8, 0]]]))
+  tests[-2, -1] = (np.array([[[4, 8, 7],
+                              [7, 8, 4]],
+                             [[1, 5, 9],
+                              [4, 8, 7]]]),
+                   np.array([[[0, 0, 0],
+                              [4, 0, 0],
+                              [7, 8, 0],
+                              [0, 8, 7],
+                              [0, 0, 4]],
+                             [[0, 0, 0],
+                              [1, 0, 0],
+                              [4, 5, 0],
+                              [0, 8, 9],
+                              [0, 0, 7]]]))
+  tests[-2, 1] = (np.array([[[2, 6, 0],
+                             [1, 5, 9],
+                             [4, 8, 7],
+                             [7, 8, 4]],
+                            [[2, 3, 0],
+                             [3, 2, 6],
+                             [1, 5, 9],
+                             [4, 8, 7]]]),
+                  np.array([[[1, 2, 0],
+                             [4, 5, 6],
+                             [7, 8, 9],
+                             [0, 8, 7],
+                             [0, 0, 4]],
+                            [[3, 2, 0],
+                             [1, 2, 3],
+                             [4, 5, 6],
+                             [0, 8, 9],
+                             [0, 0, 7]]]))
+  tests[1, 2] = (np.array([[[3, 0],
+                            [2, 6]],
+                           [[1, 0],
+                            [2, 3]]]),
+                 np.array([[[0, 2, 3],
+                            [0, 0, 6],
+                            [0, 0, 0],
+                            [0, 0, 0],
+                            [0, 0, 0]],
+                           [[0, 2, 1],
+                            [0, 0, 3],
+                            [0, 0, 0],
+                            [0, 0, 0],
+                            [0, 0, 0]]]))
+  # pyformat: enable
+  return (mat, repack_diagonals_in_tests(tests, align))
+
+
+def fat_cases(align=None):
+  # pyformat: disable
+  mat = np.array([[[1, 2, 3, 4],
+                   [5, 6, 7, 8],
+                   [9, 1, 2, 3]],
+                  [[4, 5, 6, 7],
+                   [8, 9, 1, 2],
+                   [3, 4, 5, 6]]])
+  tests = dict()
+  # tests[d_lower, d_upper] = (compact_diagonals, padded_diagonals)
+  tests[0, 0] = (np.array([[1, 6, 2],
+                           [4, 9, 5]]),
+                 np.array([[[1, 0, 0, 0],
+                            [0, 6, 0, 0],
+                            [0, 0, 2, 0]],
+                           [[4, 0, 0, 0],
+                            [0, 9, 0, 0],
+                            [0, 0, 5, 0]]]))
+  tests[2, 2] = (np.array([[3, 8],
+                           [6, 2]]),
+                 np.array([[[0, 0, 3, 0],
+                            [0, 0, 0, 8],
+                            [0, 0, 0, 0]],
+                           [[0, 0, 6, 0],
+                            [0, 0, 0, 2],
+                            [0, 0, 0, 0]]]))
+  tests[-2, 0] = (np.array([[[1, 6, 2],
+                             [5, 1, 0],
+                             [9, 0, 0]],
+                            [[4, 9, 5],
+                             [8, 4, 0],
+                             [3, 0, 0]]]),
+                  np.array([[[1, 0, 0, 0],
+                             [5, 6, 0, 0],
+                             [9, 1, 2, 0]],
+                            [[4, 0, 0, 0],
+                             [8, 9, 0, 0],
+                             [3, 4, 5, 0]]]))
+  tests[-1, 1] = (np.array([[[2, 7, 3],
+                             [1, 6, 2],
+                             [5, 1, 0]],
+                            [[5, 1, 6],
+                             [4, 9, 5],
+                             [8, 4, 0]]]),
+                  np.array([[[1, 2, 0, 0],
+                             [5, 6, 7, 0],
+                             [0, 1, 2, 3]],
+                            [[4, 5, 0, 0],
+                             [8, 9, 1, 0],
+                             [0, 4, 5, 6]]]))
+  tests[0, 3] = (np.array([[[4, 0, 0],
+                            [3, 8, 0],
+                            [2, 7, 3],
+                            [1, 6, 2]],
+                           [[7, 0, 0],
+                            [6, 2, 0],
+                            [5, 1, 6],
+                            [4, 9, 5]]]),
+                 np.array([[[1, 2, 3, 4],
+                            [0, 6, 7, 8],
+                            [0, 0, 2, 3]],
+                           [[4, 5, 6, 7],
+                            [0, 9, 1, 2],
+                            [0, 0, 5, 6]]]))
+  # pyformat: enable
+  return (mat, repack_diagonals_in_tests(tests, align))
+
+
+def all_tests(align=None):
+  return [square_cases(align), tall_cases(align), fat_cases(align)]
+
+
+class MatrixDiagTest(xla_test.XLATestCase):
+
+  def _assertOpOutputMatchesExpected(self,
+                                     params,
+                                     solution,
+                                     high_level=True,
+                                     rtol=1e-3,
+                                     atol=1e-5):
+    """Verifies that matrix_diag produces `solution` when fed `params`.
+
+    Args:
+      params: dictionary containing input parameters to matrix_diag.
+      solution: numpy array representing the expected output of matrix_diag.
+      high_level: call high_level matrix_diag
+      rtol: relative tolerance for equality test.
+      atol: absolute tolerance for equality test.
+    """
+    diagonal = params["diagonal"]
+    with self.session() as session:
+      for dtype in self.numeric_types - {np.int8, np.uint8}:
+        expected = solution.astype(dtype)
+        with self.test_scope():
+          params["diagonal"] = array_ops.placeholder(
+              dtype, diagonal.shape, name="diagonal")
+          if high_level:
+            # wraps gen_array_ops.matrix_diag_v3
+            output = array_ops.matrix_diag(**params)
+          else:
+            # TODO(b/201086188): Remove this case once MatrixDiag V1 is removed.
+            output = gen_array_ops.matrix_diag(**params)
+        result = session.run(output,
+                             {params["diagonal"]: diagonal.astype(dtype)})
+        self.assertEqual(output.dtype, expected.dtype)
+        self.assertAllCloseAccordingToType(
+            expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03)
+
+  # Generic tests applicable to both v1 and v2 ops.
+  # Originally from unary_ops_tests.py.
+  def _testV1Level(self, high_level):
+    # pyformat: disable
+    vecs1 = np.array([[1, 2],
+                      [3, 4]])
+    solution1 = np.array([[[1, 0], [0, 2]],
+                          [[3, 0], [0, 4]]])
+    vecs2 = np.array([1, 2, 3, 4])
+    solution2 = np.array([[1, 0, 0, 0],
+                          [0, 2, 0, 0],
+                          [0, 0, 3, 0],
+                          [0, 0, 0, 4]])
+    vecs3 = np.array([[[1, 2, 3],
+                       [4, 5, 6]],
+                      [[7,  8,  9],  # pylint: disable=bad-whitespace
+                       [10, 11, 12]]])
+    solution3 = np.array([[[[1, 0, 0],
+                            [0, 2, 0],
+                            [0, 0, 3]],
+                           [[4, 0, 0],
+                            [0, 5, 0],
+                            [0, 0, 6]]],
+                          [[[7, 0, 0],
+                            [0, 8, 0],
+                            [0, 0, 9]],
+                           [[10, 0, 0],
+                            [0, 11, 0],
+                            [0, 0, 12]]]])
+    # pyformat: enable
+    self._assertOpOutputMatchesExpected({"diagonal": vecs1}, solution1,
+                                        high_level)
+    self._assertOpOutputMatchesExpected({"diagonal": vecs2}, solution2,
+                                        high_level)
+    self._assertOpOutputMatchesExpected({"diagonal": vecs3}, solution3,
+                                        high_level)
+
+  def testV1(self):
+    self._testV1Level(True)
+
+  def testV1LowLevel(self):
+    self._testV1Level(False)
+
+  # From here onwards are v2-only tests.
+  def testSquare(self):
+    for align in alignment_list:
+      for _, tests in [square_cases(align)]:
+        for diag_index, (vecs, solution) in tests.items():
+          params = {"diagonal": vecs[0], "k": diag_index, "align": align}
+          self._assertOpOutputMatchesExpected(params, solution[0])
+
+  def testSquareBatch(self):
+    for align in alignment_list:
+      for _, tests in [square_cases(align)]:
+        for diag_index, (vecs, solution) in tests.items():
+          params = {"diagonal": vecs, "k": diag_index, "align": align}
+          self._assertOpOutputMatchesExpected(params, solution)
+
+  def testRectangularBatch(self):
+    # Stores expected num_rows and num_cols (when the other is given).
+    # expected[(d_lower, d_upper)] = (expected_num_rows, expected_num_cols)
+    test_list = list()
+
+    # Do not align the test cases here. Re-alignment needs to happen after the
+    # solution shape is updated.
+    # Square cases:
+    expected = {
+        (-1, -1): (5, 4),
+        (-4, -3): (5, 2),
+        (-2, 1): (5, 5),
+        (2, 4): (3, 5),
+    }
+    test_list.append((expected, square_cases()))
+
+    # Tall cases
+    expected = {
+        (0, 0): (3, 3),
+        (-4, -3): (5, 2),
+        (-2, -1): (4, 3),
+        (-2, 1): (3, 3),
+        (1, 2): (2, 3)
+    }
+    test_list.append((expected, tall_cases()))
+
+    # Fat cases
+    expected = {
+        (2, 2): (2, 4),
+        (-2, 0): (3, 3),
+        (-1, 1): (3, 3),
+        (0, 3): (3, 3)
+    }
+    test_list.append((expected, fat_cases()))
+
+    # Giving both num_rows and num_cols
+    align = alignment_list[0]
+    for _, tests in [tall_cases(align), fat_cases(align)]:
+      for diag_index, (vecs, solution) in tests.items():
+        self._assertOpOutputMatchesExpected(
+            {
+                "diagonal": vecs,
+                "k": diag_index,
+                "num_rows": solution.shape[-2],
+                "num_cols": solution.shape[-1],
+                "align": align
+            }, solution)
+
+    # We go through each alignment in a round-robin manner.
+    align_index = 0
+
+    # Giving just num_rows or num_cols.
+    for expected, (_, tests) in test_list:
+      for diag_index, (new_num_rows, new_num_cols) in expected.items():
+        align = alignment_list[align_index]
+        align_index = (align_index + 1) % len(alignment_list)
+        vecs, solution = tests[diag_index]
+        solution_given_num_rows = solution.take(
+            indices=range(new_num_cols), axis=-1)
+        # Repacks the diagonal input according to the new solution shape.
+        vecs_given_num_rows = repack_diagonals(
+            vecs,
+            diag_index,
+            solution_given_num_rows.shape[-2],
+            new_num_cols,
+            align=align)
+        self._assertOpOutputMatchesExpected(
+            {
+                "diagonal": vecs_given_num_rows,
+                "k": diag_index,
+                "num_rows": solution_given_num_rows.shape[-2],
+                "align": align
+            }, solution_given_num_rows)
+        solution_given_num_cols = solution.take(
+            indices=range(new_num_rows), axis=-2)
+        # Repacks the diagonal input according to the new solution shape.
+        vecs_given_num_cols = repack_diagonals(
+            vecs,
+            diag_index,
+            new_num_rows,
+            solution_given_num_cols.shape[-1],
+            align=align)
+        self._assertOpOutputMatchesExpected(
+            {
+                "diagonal": vecs_given_num_cols,
+                "k": diag_index,
+                "num_cols": solution_given_num_cols.shape[-1],
+                "align": align
+            }, solution_given_num_cols)
+
+  def testPadding(self):
+    for padding_value, align in zip_to_first_list_length([555, -11],
+                                                         alignment_list):
+      for _, tests in all_tests(align):
+        for diag_index, (vecs, solution) in tests.items():
+          mask = (solution == 0)
+          solution = solution + (mask * padding_value)
+          self._assertOpOutputMatchesExpected(
+              {
+                  "diagonal": vecs,
+                  "k": diag_index,
+                  "num_rows": solution.shape[-2],
+                  "num_cols": solution.shape[-1],
+                  "padding_value": padding_value,
+                  "align": align
+              }, solution)
+
+
+class MatrixSetDiagTest(xla_test.XLATestCase):
+
+  def _assertOpOutputMatchesExpected(self,
+                                     params,
+                                     solution,
+                                     high_level=True,
+                                     rtol=1e-3,
+                                     atol=1e-5):
+    """Verifies that matrix_set_diag produces `solution` when fed `params`.
+
+    Args:
+      params: dictionary containing input parameters to matrix_set_diag.
+      solution: numpy array representing the expected output of matrix_set_diag.
+      high_level: call high_level matrix_set_diag
+      rtol: relative tolerance for equality test.
+      atol: absolute tolerance for equality test.
+    """
+    input = params["input"]  # pylint: disable=redefined-builtin
+    diagonal = params["diagonal"]
+    with self.session() as session:
+      for dtype in self.numeric_types - {np.int8, np.uint8}:
+        expected = solution.astype(dtype)
+        with self.test_scope():
+          params["input"] = array_ops.placeholder(
+              dtype, input.shape, name="input")
+          params["diagonal"] = array_ops.placeholder(
+              dtype, diagonal.shape, name="diagonal")
+          if high_level:
+            # wraps gen_array_ops.matrix_set_diag_v3
+            output = array_ops.matrix_set_diag(**params)
+          else:
+            # TODO(b/201086188): Remove this case once MatrixDiag V1 is removed.
+            output = gen_array_ops.matrix_set_diag(**params)
+        result = session.run(
+            output, {
+                params["input"]: input.astype(dtype),
+                params["diagonal"]: diagonal.astype(dtype)
+            })
+        self.assertEqual(output.dtype, expected.dtype)
+        self.assertAllCloseAccordingToType(
+            expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03)
+
+  # Generic tests applicable to both v1 and v2 ops.
+  # Originally from binary_ops_tests.py.
+  def _testV1Level(self, high_level):
+    test_cases = list()
+
+    # pyformat: disable
+    # pylint: disable=bad-whitespace
+    # Square cases.
+    input = np.array([[0, 1, 0],  # pylint: disable=redefined-builtin
+                      [1, 0, 1],
+                      [1, 1, 1]])
+    diag = np.array([1, 2, 3])
+    solution = np.array([[1, 1, 0],
+                         [1, 2, 1],
+                         [1, 1, 3]])
+    test_cases.append(({"input": input, "diagonal": diag}, solution))
+
+    input = np.array([[[1, 0, 3],
+                       [0, 2, 0],
+                       [1, 0, 3]],
+                      [[4, 0, 4],
+                       [0, 5, 0],
+                       [2, 0, 6]]])
+    diag = np.array([[-1,  0, -3],
+                     [-4, -5, -6]])
+    solution = np.array([[[-1, 0,  3],
+                          [ 0, 0,  0],
+                          [ 1, 0, -3]],
+                         [[-4,  0,  4],
+                          [ 0, -5,  0],
+                          [ 2,  0, -6]]])
+    test_cases.append(({"input": input, "diagonal": diag}, solution))
+
+    # Rectangular cases.
+    input = np.array([[0, 1, 0],
+                      [1, 0, 1]])
+    diag = np.array([3, 4])
+    solution = np.array([[3, 1, 0],
+                         [1, 4, 1]])
+    test_cases.append(({"input": input, "diagonal": diag}, solution))
+
+    input = np.array([[0, 1],
+                      [1, 0],
+                      [1, 1]])
+    diag = np.array([3, 4])
+    solution = np.array([[3, 1],
+                         [1, 4],
+                         [1, 1]])
+    test_cases.append(({"input": input, "diagonal": diag}, solution))
+
+    input = np.array([[[1, 0, 3],
+                       [0, 2, 0]],
+                      [[4, 0, 4],
+                       [0, 5, 0]]])
+    diag = np.array([[-1, -2], [-4, -5]])
+    solution = np.array([[[-1,  0, 3],
+                          [ 0, -2, 0]],
+                         [[-4,  0, 4],
+                          [ 0, -5, 0]]])
+    test_cases.append(({"input": input, "diagonal": diag}, solution))
+    # pylint: enable=bad-whitespace
+    # pyformat: enable
+
+    for test in test_cases:
+      self._assertOpOutputMatchesExpected(test[0], test[1], high_level)
+
+  def testV1(self):
+    self._testV1Level(True)
+
+  def testV1LowLevel(self):
+    self._testV1Level(False)
+
+  # From here onwards are v2-only tests.
+  def testSingleMatrix(self):
+    for align in alignment_list:
+      for _, tests in all_tests(align):
+        for diag_index, (vecs, banded_mat) in tests.items():
+          mask = (banded_mat[0] == 0)
+          input_mat = np.random.randint(10, size=mask.shape)
+          solution = input_mat * mask + banded_mat[0]
+          self._assertOpOutputMatchesExpected(
+              {
+                  "input": input_mat,
+                  "diagonal": vecs[0],
+                  "k": diag_index,
+                  "align": align
+              }, solution)
+
+  def testBatch(self):
+    for align in alignment_list:
+      for _, tests in all_tests(align):
+        for diag_index, (vecs, banded_mat) in tests.items():
+          mask = (banded_mat == 0)
+          input_mat = np.random.randint(10, size=mask.shape)
+          solution = input_mat * mask + banded_mat
+          self._assertOpOutputMatchesExpected(
+              {
+                  "input": input_mat,
+                  "diagonal": vecs,
+                  "k": diag_index,
+                  "align": align
+              }, solution)
+
+
+class MatrixDiagPartTest(xla_test.XLATestCase):
+
+  def _assertOpOutputMatchesExpected(self,
+                                     params,
+                                     solution,
+                                     high_level=True,
+                                     rtol=1e-3,
+                                     atol=1e-5):
+    """Verifies that matrix_diag_part produces `solution` when fed `params`.
+
+    Args:
+      params: dictionary containing input parameters to matrix_diag_part.
+      solution: numpy array representing the expected output.
+      high_level: call high_level matrix_set_diag
+      rtol: relative tolerance for equality test.
+      atol: absolute tolerance for equality test.
+    """
+    input = params["input"]  # pylint: disable=redefined-builtin
+    with self.session() as session:
+      for dtype in self.numeric_types - {np.int8, np.uint8}:
+        expected = solution.astype(dtype)
+        with self.test_scope():
+          params["input"] = array_ops.placeholder(
+              dtype, input.shape, name="input")
+          if high_level:
+            # wraps gen_array_ops.matrix_diag_part_v3
+            output = array_ops.matrix_diag_part(**params)
+          else:
+            # TODO(b/201086188): Remove this case once MatrixDiag V1 is removed.
+            output = gen_array_ops.matrix_diag_part(**params)
+          output = array_ops.matrix_diag_part(**params)
+        result = session.run(output, {
+            params["input"]: input.astype(dtype),
+        })
+        self.assertEqual(output.dtype, expected.dtype)
+        self.assertAllCloseAccordingToType(
+            expected, result, rtol=rtol, atol=atol, bfloat16_rtol=0.03)
+
+  # Generic tests applicable to both v1 and v2 ops.
+  # Originally from unary_ops_tests.py.
+  def _testV1Level(self, high_level):
+    matrices = np.arange(3 * 2 * 4).reshape([3, 2, 4])
+    solution = np.array([[0, 5], [8, 13], [16, 21]])
+    self._assertOpOutputMatchesExpected({"input": matrices}, solution,
+                                        high_level)
+
+  def testV1(self):
+    self._testV1Level(True)
+
+  def testV1LowLevel(self):
+    self._testV1Level(False)
+
+  # From here onwards are v2-only tests.
+  def testSingleMatrix(self):
+    for align in alignment_list:
+      test_list = [square_cases(align), tall_cases(align), fat_cases(align)]
+      for mat, tests in test_list:
+        for diag_index, (solution, _) in tests.items():
+          self._assertOpOutputMatchesExpected(
+              {
+                  "input": mat[0],
+                  "k": diag_index,
+                  "align": align
+              }, solution[0])
+
+  def testBatch(self):
+    for align in alignment_list:
+      for mat, tests in all_tests(align):
+        for diag_index, (solution, _) in tests.items():
+          self._assertOpOutputMatchesExpected(
+              {
+                  "input": mat,
+                  "k": diag_index,
+                  "align": align
+              }, solution)
+
+  def testPadding(self):
+    for padding_value, align in zip_to_first_list_length([555, -11],
+                                                         alignment_list):
+      for mat, tests in all_tests(align):
+        for diag_index, (solution, _) in tests.items():
+          mask = (solution == 0)
+          solution = solution + (mask * padding_value)
+          self._assertOpOutputMatchesExpected(
+              {
+                  "input": mat,
+                  "k": diag_index,
+                  "padding_value": padding_value,
+                  "align": align
+              }, solution)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/matrix_inverse_op_test.py

@@ -0,0 +1,82 @@
+# Copyright 2015 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 tensorflow.ops.math_ops.matrix_inverse."""
+
+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 linalg_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class InverseOpTest(xla_test.XLATestCase):
+
+  def _verifyInverse(self, x, np_type):
+    for adjoint in False, True:
+      y = x.astype(np_type)
+      with self.session() as sess:
+        # Verify that x^{-1} * x == Identity matrix.
+        p = array_ops.placeholder(dtypes.as_dtype(y.dtype), y.shape, name="x")
+        with self.test_scope():
+          inv = linalg_ops.matrix_inverse(p, adjoint=adjoint)
+          tf_ans = math_ops.matmul(inv, p, adjoint_b=adjoint)
+          np_ans = np.identity(y.shape[-1])
+          if x.ndim > 2:
+            tiling = list(y.shape)
+            tiling[-2:] = [1, 1]
+            np_ans = np.tile(np_ans, tiling)
+        out = sess.run(tf_ans, feed_dict={p: y})
+        self.assertAllClose(np_ans, out, rtol=1e-3, atol=1e-3)
+        self.assertShapeEqual(y, tf_ans)
+
+  def _verifyInverseReal(self, x):
+    for np_type in self.float_types & {np.float64, np.float32}:
+      self._verifyInverse(x, np_type)
+
+  def _makeBatch(self, matrix1, matrix2):
+    matrix_batch = np.concatenate(
+        [np.expand_dims(matrix1, 0),
+         np.expand_dims(matrix2, 0)])
+    matrix_batch = np.tile(matrix_batch, [2, 3, 1, 1])
+    return matrix_batch
+
+  def testNonsymmetric(self):
+    # 2x2 matrices
+    matrix1 = np.array([[1., 2.], [3., 4.]])
+    matrix2 = np.array([[1., 3.], [3., 5.]])
+    self._verifyInverseReal(matrix1)
+    self._verifyInverseReal(matrix2)
+    # A multidimensional batch of 2x2 matrices
+    self._verifyInverseReal(self._makeBatch(matrix1, matrix2))
+
+  def testSymmetricPositiveDefinite(self):
+    # 2x2 matrices
+    matrix1 = np.array([[2., 1.], [1., 2.]])
+    matrix2 = np.array([[3., -1.], [-1., 3.]])
+    self._verifyInverseReal(matrix1)
+    self._verifyInverseReal(matrix2)
+    # A multidimensional batch of 2x2 matrices
+    self._verifyInverseReal(self._makeBatch(matrix1, matrix2))
+
+  def testEmpty(self):
+    self._verifyInverseReal(np.empty([0, 2, 2]))
+    self._verifyInverseReal(np.empty([2, 0, 0]))
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/matrix_solve_op_test.py

@@ -0,0 +1,81 @@
+# Copyright 2015 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 implementation of tf.linalg.solve."""
+
+import os
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.platform import googletest
+from tensorflow.python.platform import sysconfig
+
+class MatrixSolveOpTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _verifySolve(self, x, y, adjoint):
+    for np_type in self.float_types & {np.float32, np.float64}:
+      tol = 1e-4 if np_type == np.float32 else 1e-12
+      a = x.astype(np_type)
+      b = y.astype(np_type)
+      np_ans = np.linalg.solve(np.swapaxes(a, -2, -1) if adjoint else a, b)
+      with self.session() as sess:
+        with self.test_scope():
+          tf_ans = linalg_ops.matrix_solve(a, b, adjoint=adjoint)
+        out = sess.run(tf_ans)
+        self.assertEqual(tf_ans.shape, out.shape)
+        self.assertEqual(np_ans.shape, out.shape)
+        self.assertAllClose(np_ans, out, atol=tol, rtol=tol)
+
+  @parameterized.named_parameters(
+      ("Scalar", 1, 1, [], [], False),
+      ("Vector", 5, 1, [], [], False),
+      ("MultipleRHS", 5, 4, [], [], False),
+      ("Adjoint", 5, 4, [], [], True),
+      ("BatchedScalar", 1, 4, [2], [2], False),
+      ("BatchedVector", 5, 4, [2], [2], False),
+      ("BatchedRank2", 5, 4, [7, 4], [7, 4], False),
+      ("BatchedAdjoint", 5, 4, [7, 4], [7, 4], True),
+  )
+  def testSolve(self, n, nrhs, batch_dims, rhs_batch_dims, adjoint):
+    matrix = np.random.normal(-5.0, 5.0, batch_dims + [n, n])
+    rhs = np.random.normal(-5.0, 5.0, rhs_batch_dims + [n, nrhs])
+    self._verifySolve(matrix, rhs, adjoint=adjoint)
+
+  @parameterized.named_parameters(
+      ("Simple", False),
+      ("Adjoint", True),
+  )
+  def testConcurrent(self, adjoint):
+    with self.session() as sess:
+      lhs1 = random_ops.random_normal([3, 3], seed=42)
+      lhs2 = random_ops.random_normal([3, 3], seed=42)
+      rhs1 = random_ops.random_normal([3, 3], seed=42)
+      rhs2 = random_ops.random_normal([3, 3], seed=42)
+      with self.test_scope():
+        s1 = linalg_ops.matrix_solve(lhs1, rhs1, adjoint=adjoint)
+        s2 = linalg_ops.matrix_solve(lhs2, rhs2, adjoint=adjoint)
+      self.assertAllEqual(*sess.run([s1, s2]))
+
+
+if __name__ == "__main__":
+  sys_details = sysconfig.get_build_info()
+  if sys_details["is_cuda_build"]:
+    os.environ["XLA_FLAGS"] = (
+        "--xla_gpu_enable_cublaslt=true " + os.environ.get("XLA_FLAGS", "")
+    )
+  googletest.main()

tensorflow/compiler/tests/matrix_triangular_solve_op_test.py

@@ -0,0 +1,177 @@
+# Copyright 2018 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 tensorflow.ops.tf.MatrixTriangularSolve."""
+
+import itertools
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.platform import test
+
+
+def MakePlaceholder(x, dtype=None):
+  return array_ops.placeholder(
+      dtypes.as_dtype(x.dtype) if dtype is None else dtype, shape=x.shape)
+
+
+class MatrixTriangularSolveOpTest(xla_test.XLATestCase):
+
+  #  MatrixTriangularSolve defined for float64, float32, complex64, complex128
+  # (https://www.tensorflow.org/api_docs/python/tf/matrix_triangular_solve)
+  @property
+  def float_types(self):
+    return set(super(MatrixTriangularSolveOpTest,
+                     self).float_types).intersection(
+                         (np.float64, np.float32, np.complex64, np.complex128))
+
+  def _VerifyTriangularSolveBase(self, sess, placeholder_a, placeholder_ca,
+                                 placeholder_b, a, clean_a, b, verification,
+                                 atol):
+    feed_dict = {placeholder_a: a, placeholder_ca: clean_a, placeholder_b: b}
+    verification_np = sess.run(verification, feed_dict)
+    broadcasted_shape = a.shape[:-2] + (b.shape[-2], b.shape[-1])
+    broadcasted_b = b + np.zeros(shape=broadcasted_shape, dtype=b.dtype)
+    self.assertAllClose(broadcasted_b, verification_np, atol=atol)
+
+  def _VerifyTriangularSolve(self, a, b, lower, adjoint, atol, dtype=None):
+    clean_a = np.tril(a) if lower else np.triu(a)
+    with self.session() as sess:
+      placeholder_a = MakePlaceholder(a, dtype)
+      placeholder_ca = MakePlaceholder(clean_a, dtype)
+      placeholder_b = MakePlaceholder(b, dtype)
+      with self.test_scope():
+        x = linalg_ops.matrix_triangular_solve(
+            placeholder_a, placeholder_b, lower=lower, adjoint=adjoint)
+      verification = test_util.matmul_without_tf32(
+          placeholder_ca, x, adjoint_a=adjoint)
+      self._VerifyTriangularSolveBase(sess, placeholder_a, placeholder_ca,
+                                      placeholder_b, a, clean_a, b,
+                                      verification, atol)
+
+  def _VerifyTriangularSolveCombo(self, a, b, atol=1e-4, dtype=None):
+    transp = lambda x: np.swapaxes(x, -1, -2)
+    for lower, adjoint in itertools.product([True, False], repeat=2):
+      self._VerifyTriangularSolve(
+          a if lower else transp(a), b, lower, adjoint, atol, dtype=dtype)
+
+  def testBasic(self):
+    rng = np.random.RandomState(0)
+    a = np.tril(rng.randn(5, 5))
+    b = rng.randn(5, 7)
+    for dtype in self.float_types:
+      self._VerifyTriangularSolveCombo(a.astype(dtype), b.astype(dtype))
+
+  def testBfloat16(self):
+    rng = np.random.RandomState(0)
+    a = np.tril(rng.randn(5, 5))
+    b = rng.randn(5, 7)
+    self._VerifyTriangularSolveCombo(a, b, atol=5e-2, dtype=dtypes.bfloat16)
+
+  def testBasicNotActuallyTriangular(self):
+    rng = np.random.RandomState(0)
+    a = rng.randn(5, 5)  # the `a` matrix is not lower-triangular
+    b = rng.randn(5, 7)
+    for dtype in self.float_types:
+      self._VerifyTriangularSolveCombo(a.astype(dtype), b.astype(dtype))
+
+  def testBasicComplexDtypes(self):
+
+    if xla_test.test.is_built_with_rocm():
+      # The following subtest invokes the call to "BlasTrsm"
+      # That operation is currently not supported on the ROCm platform
+      self.skipTest("BlasTrsm op for complex types is not supported in ROCm")
+
+    rng = np.random.RandomState(0)
+    a = np.tril(rng.randn(5, 5) + rng.randn(5, 5) * 1j)
+    b = rng.randn(5, 7) + rng.randn(5, 7) * 1j
+    for dtype in self.complex_types:
+      self._VerifyTriangularSolveCombo(a.astype(dtype), b.astype(dtype))
+
+  def testBatch(self):
+    rng = np.random.RandomState(0)
+    shapes = [((4, 3, 3), (4, 3, 5)), ((1, 2, 2), (1, 2, 1)),
+              ((1, 1, 1), (1, 1, 2)), ((2, 3, 4, 4), (2, 3, 4, 1))]
+    tuples = itertools.product(self.float_types, shapes)
+    for dtype, (a_shape, b_shape) in tuples:
+      n = a_shape[-1]
+      a = np.tril(rng.rand(*a_shape) - 0.5) / (2.0 * n) + np.eye(n)
+      b = rng.randn(*b_shape)
+      self._VerifyTriangularSolveCombo(
+          a.astype(dtype), b.astype(dtype), atol=1e-3)
+
+  def testBatchBroadcast(self):
+    rng = np.random.RandomState(0)
+    shapes = [((3, 3), (4, 3, 5)), ((1, 2, 2), (3, 2, 1)), ((1, 1), (1, 1, 2)),
+              ((1, 3, 4, 4), (2, 1, 4, 1))]
+    tuples = itertools.product(self.float_types, shapes)
+    for dtype, (a_shape, b_shape) in tuples:
+      n = a_shape[-1]
+      a = np.tril(rng.rand(*a_shape) - 0.5) / (2.0 * n) + np.eye(n)
+      b = rng.randn(*b_shape)
+      self._VerifyTriangularSolveCombo(
+          a.astype(dtype), b.astype(dtype), atol=1e-3)
+
+  def testLarge(self):
+    n = 1024
+    rng = np.random.RandomState(0)
+    a = np.tril(rng.rand(n, n) - 0.5) / (2.0 * n) + np.eye(n)
+    b = rng.randn(n, n)
+    self._VerifyTriangularSolve(
+        a.astype(np.float32), b.astype(np.float32), True, False, 1e-4)
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testNonSquareCoefficientMatrix(self):
+    rng = np.random.RandomState(0)
+    for dtype in self.float_types:
+      a = rng.randn(3, 4).astype(dtype)
+      b = rng.randn(4, 4).astype(dtype)
+      with self.test_scope():
+        with self.assertRaises((ValueError, errors.InvalidArgumentError)):
+          linalg_ops.matrix_triangular_solve(a, b)
+
+  @test_util.run_v2_only  # Different error types
+  @test_util.disable_mlir_bridge("Error handling")
+  def testWrongDimensionsV2(self):
+    randn = np.random.RandomState(0).randn
+    for dtype in self.float_types:
+      lhs = constant_op.constant(randn(3, 3), dtype=dtype)
+      rhs = constant_op.constant(randn(4, 3), dtype=dtype)
+      with self.assertRaises(errors.InvalidArgumentError):
+        linalg_ops.matrix_triangular_solve(lhs, rhs)
+      with self.assertRaises(errors.InvalidArgumentError):
+        linalg_ops.matrix_triangular_solve(lhs, rhs)
+
+  @test_util.run_v1_only("Different error types")
+  @test_util.disable_mlir_bridge("Error handling")
+  def testWrongDimensionsV1(self):
+    randn = np.random.RandomState(0).randn
+    for dtype in self.float_types:
+      lhs = constant_op.constant(randn(3, 3), dtype=dtype)
+      rhs = constant_op.constant(randn(4, 3), dtype=dtype)
+      with self.assertRaises(ValueError):
+        linalg_ops.matrix_triangular_solve(lhs, rhs)
+      with self.assertRaises(ValueError):
+        linalg_ops.matrix_triangular_solve(lhs, rhs)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/mean_op_test.py

@@ -0,0 +1,39 @@
+# Copyright 2023 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 Mean operators."""
+
+from tensorflow.python.eager.def_function import function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class MeanOpTest(test.TestCase):
+  def testReduceMeanOverflow(self):
+    @function(jit_compile=True)
+    def tpu_computation():
+      shape = (4, 32, 512, 512, 96)
+      dtype = dtypes.float32
+      test_ones = array_ops.ones(shape=shape, dtype=dtype)
+      return math_ops.reduce_mean(test_ones)
+    with ops.device("TPU:0"):
+      result = tpu_computation()
+      self.assertAllClose(
+          result, 1.0, rtol=1e-4, atol=1e-4)
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/momentum_test.py

@@ -0,0 +1,186 @@
+# 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.
+# ==============================================================================
+"""Tests for Momentum."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import momentum as momentum_lib
+
+
+class MomentumOptimizerTest(xla_test.XLATestCase):
+
+  def _update_nesterov_momentum_numpy(self, var, accum, g, lr, momentum):
+    var += accum * lr * momentum
+    accum = accum * momentum + g
+    var -= lr * accum
+    var -= accum * lr * momentum
+    return var, accum
+
+  def testBasic(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        mom_opt = momentum_lib.MomentumOptimizer(
+            learning_rate=2.0, momentum=0.9)
+        mom_update = mom_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Check we have slots
+        self.assertEqual(["momentum"], mom_opt.get_slot_names())
+        slot0 = mom_opt.get_slot(var0, "momentum")
+        self.assertEqual(slot0.get_shape(), var0.get_shape())
+        self.assertFalse(slot0 in variables.trainable_variables())
+        slot1 = mom_opt.get_slot(var1, "momentum")
+        self.assertEqual(slot1.get_shape(), var1.get_shape())
+        self.assertFalse(slot1 in variables.trainable_variables())
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Step 1: the momentum accumulators where 0. So we should see a normal
+        # update: v -= grad * learning_rate
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([0.1, 0.1]), self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.01, 0.01]), self.evaluate(slot1))
+        # Check that the parameters have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]),
+            self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]),
+            self.evaluate(var1))
+        # Step 2: the momentum accumulators contain the previous update.
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]),
+            self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 0.01)]),
+            self.evaluate(slot1))
+        # Check that the parameters have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0),
+                2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0)
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                2.98 - ((0.9 * 0.01 + 0.01) * 2.0),
+                3.98 - ((0.9 * 0.01 + 0.01) * 2.0)
+            ]), self.evaluate(var1))
+
+  def testNesterovMomentum(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([0.1, 0.2], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([0.3, 0.4], dtype=dtype)
+        var0_np = np.array([0.1, 0.2], dtype=dtype)
+        var1_np = np.array([0.3, 0.4], dtype=dtype)
+        accum0_np = np.array([0.0, 0.0], dtype=dtype)
+        accum1_np = np.array([0.0, 0.0], dtype=dtype)
+        cost = 0.4 * var0 * var0 + 0.9 * var1
+        global_step = resource_variable_ops.ResourceVariable(
+            array_ops.zeros([], dtypes.int32), name="global_step")
+        mom_op = momentum_lib.MomentumOptimizer(
+            learning_rate=0.1, momentum=0.9, use_nesterov=True)
+        opt_op = mom_op.minimize(cost, global_step, [var0, var1])
+        self.evaluate(variables.global_variables_initializer())
+        for _ in range(1, 5):
+          opt_op.run()
+          var0_np, accum0_np = self._update_nesterov_momentum_numpy(
+              var0_np, accum0_np, var0_np * 0.8, 0.1, 0.9)
+          var1_np, accum1_np = self._update_nesterov_momentum_numpy(
+              var1_np, accum1_np, 0.9, 0.1, 0.9)
+          self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+          self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+  def testTensorLearningRateAndMomentum(self):
+    for dtype in self.float_types:
+      with self.session(), self.test_scope():
+        var0 = resource_variable_ops.ResourceVariable([1.0, 2.0], dtype=dtype)
+        var1 = resource_variable_ops.ResourceVariable([3.0, 4.0], dtype=dtype)
+        grads0 = constant_op.constant([0.1, 0.1], dtype=dtype)
+        grads1 = constant_op.constant([0.01, 0.01], dtype=dtype)
+        mom_opt = momentum_lib.MomentumOptimizer(
+            learning_rate=constant_op.constant(2.0),
+            momentum=constant_op.constant(0.9))
+        mom_update = mom_opt.apply_gradients(
+            zip([grads0, grads1], [var0, var1]))
+        self.evaluate(variables.global_variables_initializer())
+        # Check we have slots
+        self.assertEqual(["momentum"], mom_opt.get_slot_names())
+        slot0 = mom_opt.get_slot(var0, "momentum")
+        self.assertEqual(slot0.get_shape(), var0.get_shape())
+        self.assertFalse(slot0 in variables.trainable_variables())
+        slot1 = mom_opt.get_slot(var1, "momentum")
+        self.assertEqual(slot1.get_shape(), var1.get_shape())
+        self.assertFalse(slot1 in variables.trainable_variables())
+
+        # Fetch params to validate initial values
+        self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+        self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+        # Step 1: the momentum accumulators where 0. So we should see a normal
+        # update: v -= grad * learning_rate
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([0.1, 0.1]), self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([0.01, 0.01]), self.evaluate(slot1))
+        # Check that the parameters have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([1.0 - (0.1 * 2.0), 2.0 - (0.1 * 2.0)]),
+            self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([3.0 - (0.01 * 2.0), 4.0 - (0.01 * 2.0)]),
+            self.evaluate(var1))
+        # Step 2: the momentum accumulators contain the previous update.
+        mom_update.run()
+        # Check that the momentum accumulators have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.1 + 0.1), (0.9 * 0.1 + 0.1)]),
+            self.evaluate(slot0))
+        self.assertAllCloseAccordingToType(
+            np.array([(0.9 * 0.01 + 0.01), (0.9 * 0.01 + 0.01)]),
+            self.evaluate(slot1))
+        # Check that the parameters have been updated.
+        self.assertAllCloseAccordingToType(
+            np.array([
+                1.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0),
+                2.0 - (0.1 * 2.0) - ((0.9 * 0.1 + 0.1) * 2.0)
+            ]), self.evaluate(var0))
+        self.assertAllCloseAccordingToType(
+            np.array([
+                2.98 - ((0.9 * 0.01 + 0.01) * 2.0),
+                3.98 - ((0.9 * 0.01 + 0.01) * 2.0)
+            ]), self.evaluate(var1))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/nary_ops_test.py

@@ -0,0 +1,291 @@
+# 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.
+# ==============================================================================
+"""Test cases for operators with > 3 or arbitrary numbers of arguments."""
+
+import unittest
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class NAryOpsTest(xla_test.XLATestCase):
+
+  def _testNAry(self, op, args, expected, equality_fn=None):
+    with self.session() as session:
+      with self.test_scope():
+        placeholders = [
+            array_ops.placeholder(dtypes.as_dtype(arg.dtype), arg.shape)
+            for arg in args
+        ]
+        feeds = {placeholders[i]: args[i] for i in range(0, len(args))}
+        output = op(placeholders)
+      result = session.run(output, feeds)
+      if not equality_fn:
+        equality_fn = self.assertAllClose
+      equality_fn(result, expected, rtol=1e-3)
+
+  def _nAryListCheck(self, results, expected, **kwargs):
+    self.assertEqual(len(results), len(expected))
+    for (r, e) in zip(results, expected):
+      self.assertAllClose(r, e, **kwargs)
+
+  def _testNAryLists(self, op, args, expected):
+    self._testNAry(op, args, expected, equality_fn=self._nAryListCheck)
+
+  def testFloat(self):
+    self._testNAry(math_ops.add_n,
+                   [np.array([[1, 2, 3]], dtype=np.float32)],
+                   expected=np.array([[1, 2, 3]], dtype=np.float32))
+
+    self._testNAry(math_ops.add_n,
+                   [np.array([1, 2], dtype=np.float32),
+                    np.array([10, 20], dtype=np.float32)],
+                   expected=np.array([11, 22], dtype=np.float32))
+    self._testNAry(math_ops.add_n,
+                   [np.array([-4], dtype=np.float32),
+                    np.array([10], dtype=np.float32),
+                    np.array([42], dtype=np.float32)],
+                   expected=np.array([48], dtype=np.float32))
+
+  def testComplex(self):
+    for dtype in self.complex_types:
+      self._testNAry(
+          math_ops.add_n, [np.array([[1 + 2j, 2 - 3j, 3 + 4j]], dtype=dtype)],
+          expected=np.array([[1 + 2j, 2 - 3j, 3 + 4j]], dtype=dtype))
+
+      self._testNAry(
+          math_ops.add_n, [
+              np.array([1 + 2j, 2 - 3j], dtype=dtype),
+              np.array([10j, 20], dtype=dtype)
+          ],
+          expected=np.array([1 + 12j, 22 - 3j], dtype=dtype))
+      self._testNAry(
+          math_ops.add_n, [
+              np.array([-4, 5j], dtype=dtype),
+              np.array([2 + 10j, -2], dtype=dtype),
+              np.array([42j, 3 + 3j], dtype=dtype)
+          ],
+          expected=np.array([-2 + 52j, 1 + 8j], dtype=dtype))
+
+  @unittest.skip("IdentityN is temporarily CompilationOnly as workaround")
+  def testIdentityN(self):
+    self._testNAryLists(array_ops.identity_n,
+                        [np.array([[1, 2, 3]], dtype=np.float32)],
+                        expected=[np.array([[1, 2, 3]], dtype=np.float32)])
+    self._testNAryLists(array_ops.identity_n,
+                        [np.array([[1, 2], [3, 4]], dtype=np.float32),
+                         np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)],
+                        expected=[
+                            np.array([[1, 2], [3, 4]], dtype=np.float32),
+                            np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)])
+    self._testNAryLists(array_ops.identity_n,
+                        [np.array([[1], [2], [3], [4]], dtype=np.int32),
+                         np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)],
+                        expected=[
+                            np.array([[1], [2], [3], [4]], dtype=np.int32),
+                            np.array([[3, 2, 1], [6, 5, 1]], dtype=np.float32)])
+
+  def testConcat(self):
+    self._testNAry(
+        lambda x: array_ops.concat(x, 0), [
+            np.array(
+                [[1, 2, 3], [4, 5, 6]], dtype=np.float32), np.array(
+                    [[7, 8, 9], [10, 11, 12]], dtype=np.float32)
+        ],
+        expected=np.array(
+            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=np.float32))
+
+    self._testNAry(
+        lambda x: array_ops.concat(x, 1), [
+            np.array(
+                [[1, 2, 3], [4, 5, 6]], dtype=np.float32), np.array(
+                    [[7, 8, 9], [10, 11, 12]], dtype=np.float32)
+        ],
+        expected=np.array(
+            [[1, 2, 3, 7, 8, 9], [4, 5, 6, 10, 11, 12]], dtype=np.float32))
+
+  def testOneHot(self):
+    with self.session() as session, self.test_scope():
+      indices = array_ops.constant(np.array([[2, 3], [0, 1]], dtype=np.int32))
+      op = array_ops.one_hot(indices,
+                             np.int32(4),
+                             on_value=np.float32(7), off_value=np.float32(3))
+      output = session.run(op)
+      expected = np.array([[[3, 3, 7, 3], [3, 3, 3, 7]],
+                           [[7, 3, 3, 3], [3, 7, 3, 3]]],
+                          dtype=np.float32)
+      self.assertAllEqual(output, expected)
+
+      op = array_ops.one_hot(indices,
+                             np.int32(4),
+                             on_value=np.int32(2), off_value=np.int32(1),
+                             axis=1)
+      output = session.run(op)
+      expected = np.array([[[1, 1], [1, 1], [2, 1], [1, 2]],
+                           [[2, 1], [1, 2], [1, 1], [1, 1]]],
+                          dtype=np.int32)
+      self.assertAllEqual(output, expected)
+
+  def testSplitV(self):
+    with self.session() as session:
+      with self.test_scope():
+        output = session.run(
+            array_ops.split(np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 0, 1, 2]],
+                                     dtype=np.float32),
+                            [2, 2], 1))
+        expected = [np.array([[1, 2], [5, 6], [9, 0]], dtype=np.float32),
+                    np.array([[3, 4], [7, 8], [1, 2]], dtype=np.float32)]
+        self.assertAllEqual(output, expected)
+
+  def testSplitVNegativeSizes(self):
+    with self.session() as session:
+      with self.test_scope():
+        with self.assertRaisesRegex(
+            (ValueError, errors.InvalidArgumentError),
+            "Split size at index 1 must be >= .*. Got: -2",
+        ):
+          _ = session.run(
+              array_ops.split(np.array([1, 2, 3], dtype=np.float32), [-1, -2],
+                              axis=0))
+
+  def testStridedSlice(self):
+    self._testNAry(lambda x: array_ops.strided_slice(*x),
+                   [np.array([[], [], []], dtype=np.float32),
+                    np.array([1, 0], dtype=np.int32),
+                    np.array([3, 0], dtype=np.int32),
+                    np.array([1, 1], dtype=np.int32)],
+                   expected=np.array([[], []], dtype=np.float32))
+
+    if np.int64 in self.int_types:
+      self._testNAry(
+          lambda x: array_ops.strided_slice(*x), [
+              np.array([[], [], []], dtype=np.float32), np.array(
+                  [1, 0], dtype=np.int64), np.array([3, 0], dtype=np.int64),
+              np.array([1, 1], dtype=np.int64)
+          ],
+          expected=np.array([[], []], dtype=np.float32))
+
+    self._testNAry(lambda x: array_ops.strided_slice(*x),
+                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                             dtype=np.float32),
+                    np.array([1, 1], dtype=np.int32),
+                    np.array([3, 3], dtype=np.int32),
+                    np.array([1, 1], dtype=np.int32)],
+                   expected=np.array([[5, 6], [8, 9]], dtype=np.float32))
+
+    self._testNAry(lambda x: array_ops.strided_slice(*x),
+                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                             dtype=np.float32),
+                    np.array([0, 2], dtype=np.int32),
+                    np.array([2, 0], dtype=np.int32),
+                    np.array([1, -1], dtype=np.int32)],
+                   expected=np.array([[3, 2], [6, 5]], dtype=np.float32))
+
+    self._testNAry(lambda x: x[0][0:2, array_ops.newaxis, ::-1],
+                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                             dtype=np.float32)],
+                   expected=np.array([[[3, 2, 1]], [[6, 5, 4]]],
+                                     dtype=np.float32))
+
+    self._testNAry(lambda x: x[0][1, :, array_ops.newaxis],
+                   [np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+                             dtype=np.float32)],
+                   expected=np.array([[4], [5], [6]], dtype=np.float32))
+
+  def testStridedSliceGrad(self):
+    # Tests cases where input shape is empty.
+    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
+                   [np.array([], dtype=np.int32),
+                    np.array([], dtype=np.int32),
+                    np.array([], dtype=np.int32),
+                    np.array([], dtype=np.int32),
+                    np.float32(0.5)],
+                   expected=np.array(np.float32(0.5), dtype=np.float32))
+
+    # Tests case where input shape is non-empty, but gradients are empty.
+    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
+                   [np.array([3], dtype=np.int32),
+                    np.array([0], dtype=np.int32),
+                    np.array([0], dtype=np.int32),
+                    np.array([1], dtype=np.int32),
+                    np.array([], dtype=np.float32)],
+                   expected=np.array([0, 0, 0], dtype=np.float32))
+
+    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
+                   [np.array([3, 0], dtype=np.int32),
+                    np.array([1, 0], dtype=np.int32),
+                    np.array([3, 0], dtype=np.int32),
+                    np.array([1, 1], dtype=np.int32),
+                    np.array([[], []], dtype=np.float32)],
+                   expected=np.array([[], [], []], dtype=np.float32))
+
+    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
+                   [np.array([3, 3], dtype=np.int32),
+                    np.array([1, 1], dtype=np.int32),
+                    np.array([3, 3], dtype=np.int32),
+                    np.array([1, 1], dtype=np.int32),
+                    np.array([[5, 6], [8, 9]], dtype=np.float32)],
+                   expected=np.array([[0, 0, 0], [0, 5, 6], [0, 8, 9]],
+                                     dtype=np.float32))
+
+    def ssg_test(x):
+      return array_ops.strided_slice_grad(*x, shrink_axis_mask=0x4,
+                                          new_axis_mask=0x1)
+
+    self._testNAry(ssg_test,
+                   [np.array([3, 1, 3], dtype=np.int32),
+                    np.array([0, 0, 0, 2], dtype=np.int32),
+                    np.array([0, 3, 1, -4], dtype=np.int32),
+                    np.array([1, 2, 1, -3], dtype=np.int32),
+                    np.array([[[1], [2]]], dtype=np.float32)],
+                   expected=np.array([[[0, 0, 1]], [[0, 0, 0]], [[0, 0, 2]]],
+                                     dtype=np.float32))
+
+    ssg_test2 = lambda x: array_ops.strided_slice_grad(*x, new_axis_mask=0x15)
+    self._testNAry(ssg_test2,
+                   [np.array([4, 4], dtype=np.int32),
+                    np.array([0, 0, 0, 1, 0], dtype=np.int32),
+                    np.array([0, 3, 0, 4, 0], dtype=np.int32),
+                    np.array([1, 2, 1, 2, 1], dtype=np.int32),
+                    np.array([[[[[1], [2]]], [[[3], [4]]]]], dtype=np.float32)],
+                   expected=np.array([[0, 1, 0, 2], [0, 0, 0, 0], [0, 3, 0, 4],
+                                      [0, 0, 0, 0]], dtype=np.float32))
+
+    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
+                   [np.array([3, 3], dtype=np.int32),
+                    np.array([0, 2], dtype=np.int32),
+                    np.array([2, 0], dtype=np.int32),
+                    np.array([1, -1], dtype=np.int32),
+                    np.array([[1, 2], [3, 4]], dtype=np.float32)],
+                   expected=np.array([[0, 2, 1], [0, 4, 3], [0, 0, 0]],
+                                     dtype=np.float32))
+
+    self._testNAry(lambda x: array_ops.strided_slice_grad(*x),
+                   [np.array([3, 3], dtype=np.int32),
+                    np.array([2, 2], dtype=np.int32),
+                    np.array([0, 1], dtype=np.int32),
+                    np.array([-1, -2], dtype=np.int32),
+                    np.array([[1], [2]], dtype=np.float32)],
+                   expected=np.array([[0, 0, 0], [0, 0, 2], [0, 0, 1]],
+                                     dtype=np.float32))
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/nullary_ops_test.py

@@ -0,0 +1,76 @@
+# 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.
+# ==============================================================================
+"""Test cases for operators with no arguments."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.platform import googletest
+
+
+class NullaryOpsTest(xla_test.XLATestCase):
+
+  def _testNullary(self, op, expected):
+    with self.session() as session:
+      with self.test_scope():
+        output = op()
+      result = session.run(output)
+      self.assertAllClose(result, expected, rtol=1e-3)
+
+  def testNoOp(self):
+    with self.session():
+      with self.test_scope():
+        output = control_flow_ops.no_op()
+      # This should not crash.
+      output.run()
+
+  def testConstants(self):
+    for dtype in self.numeric_types:
+      constants = [
+          dtype(42),
+          np.array([], dtype=dtype),
+          np.array([1, 2], dtype=dtype),
+          np.array([7, 7, 7, 7, 7], dtype=dtype),
+          np.array([[1, 2, 3], [4, 5, 6]], dtype=dtype),
+          np.array([[[1, 2], [3, 4], [5, 6]], [[10, 20], [30, 40], [50, 60]]],
+                   dtype=dtype),
+          np.array([[[]], [[]]], dtype=dtype),
+          np.array([[[[1]]]], dtype=dtype),
+      ]
+      for c in constants:
+        self._testNullary(lambda c=c: constant_op.constant(c), expected=c)
+
+  def testComplexConstants(self):
+    for dtype in self.complex_types:
+      constants = [
+          dtype(42 + 3j),
+          np.array([], dtype=dtype),
+          np.ones([50], dtype=dtype) * (3 + 4j),
+          np.array([1j, 2 + 1j], dtype=dtype),
+          np.array([[1, 2j, 7j], [4, 5, 6]], dtype=dtype),
+          np.array([[[1, 2], [3, 4 + 6j], [5, 6]],
+                    [[10 + 7j, 20], [30, 40], [50, 60]]],
+                   dtype=dtype),
+          np.array([[[]], [[]]], dtype=dtype),
+          np.array([[[[1 + 3j]]]], dtype=dtype),
+      ]
+      for c in constants:
+        self._testNullary(lambda c=c: constant_op.constant(c), expected=c)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/placeholder_test.py

@@ -0,0 +1,44 @@
+# Copyright 2018 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 handling of placeholder_with_default."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import googletest
+
+
+class PlaceholderTest(xla_test.XLATestCase):
+
+  def test_placeholder_with_default_default(self):
+    with self.session() as sess, self.test_scope():
+      v = resource_variable_ops.ResourceVariable(4.0)
+      ph = array_ops.placeholder_with_default(v, shape=[])
+      out = ph * 2
+      sess.run(variables.variables_initializer([v]))
+      self.assertEqual(8.0, self.evaluate(out))
+
+  def test_placeholder_with_default_fed(self):
+    with self.session() as sess, self.test_scope():
+      v = resource_variable_ops.ResourceVariable(4.0)
+      ph = array_ops.placeholder_with_default(v, shape=[])
+      out = ph * 2
+      sess.run(variables.variables_initializer([v]))
+      self.assertEqual(2.0, sess.run(out, {ph: 1.0}))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/plugin.bzl

@@ -0,0 +1,29 @@
+# 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.
+# ==============================================================================
+"""Additional XLA devices to be included in the unit test suite."""
+
+# If you wish to edit this file without checking it into the repo, consider:
+#   git update-index --assume-unchanged tensorflow/compiler/tests/plugin.bzl
+
+plugins = {
+    #"example": {
+    #  "device":"XLA_MY_DEVICE",
+    #  "types":"DT_FLOAT,DT_HALF,DT_INT32",
+    #   "tags":[],
+    #   "args":["--disabled_manifest=tensorflow/compiler/plugin/example/disabled_manifest.txt"],
+    #   "data":["//tensorflow/compiler/plugin/example:disabled_manifest.txt"],
+    #   "deps":[],
+    #},
+}

tensorflow/compiler/tests/pooling_ops_3d_test.py

@@ -0,0 +1,457 @@
+# 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 3d pooling operations."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import test
+
+
+# Wrapper around AvgPoolGrad that ignores extra arguments needed by
+# MaxPoolGrad.
+def _AvgPoolGrad(inputs, outputs, output_gradients, ksize, strides, padding):
+  del outputs  # Unused by average-pooling gradients.
+  return gen_nn_ops.avg_pool3d_grad(
+      inputs.get_shape().as_list(),
+      output_gradients,
+      ksize=ksize,
+      strides=strides,
+      padding=padding)
+
+
+class Pooling3DTest(xla_test.XLATestCase):
+
+  def _VerifyValues(self, pool_func, input_sizes, window, strides, padding,
+                    expected):
+    """Verifies the output values of the pooling function.
+
+    Args:
+      pool_func: Function to be called: co.MaxPool, co.AvgPool.
+      input_sizes: Input tensor dimensions.
+      window: Tuple of kernel dims: planes, rows, cols.
+      strides: Tuple of strides for dims: planes, rows, cols.
+      padding: Padding type.
+      expected: An array containing the expected operation outputs.
+    """
+    total_size = 1
+    for s in input_sizes:
+      total_size *= s
+    # Initializes the input tensor with array containing incrementing
+    # numbers from 1.
+    x = np.arange(1.0, total_size + 1, dtype=np.float32)
+    x = x.reshape(input_sizes)
+    with self.session() as sess, self.test_scope():
+      inputs = array_ops.placeholder(dtypes.float32)
+      t = pool_func(
+          inputs,
+          ksize=[1] + window + [1],
+          strides=[1] + strides + [1],
+          padding=padding)
+      vals = sess.run(t, {inputs: x})
+    # Verifies values.
+    actual = vals.flatten()
+    self.assertAllClose(expected, actual)
+
+  def testAvgPool3dValidPadding(self):
+    expected_output = [20.5, 21.5, 22.5]
+    self._VerifyValues(
+        nn_ops.avg_pool3d,
+        input_sizes=[1, 3, 3, 3, 3],
+        window=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="VALID",
+        expected=expected_output)
+
+  def testAvgPool3dSamePadding(self):
+    expected_output = [20.5, 21.5, 22.5, 26.5, 27.5, 28.5]
+    self._VerifyValues(
+        nn_ops.avg_pool3d,
+        input_sizes=[1, 2, 2, 4, 3],
+        window=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="SAME",
+        expected=expected_output)
+
+  def testAvgPool3dSamePaddingDifferentStrides(self):
+    expected_output = [1.5, 4.5, 7.5, 17.5, 20.5, 23.5, 33.5, 36.5, 39.5]
+    self._VerifyValues(
+        nn_ops.avg_pool3d,
+        input_sizes=[1, 5, 8, 1, 1],
+        window=[1, 2, 3],
+        strides=[2, 3, 1],
+        padding="SAME",
+        expected=expected_output)
+
+  def testMaxPool3dValidPadding(self):
+    expected_output = [40.0, 41.0, 42.0]
+    self._VerifyValues(
+        nn_ops.max_pool3d,
+        input_sizes=[1, 3, 3, 3, 3],
+        window=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="VALID",
+        expected=expected_output)
+
+  def testMaxPool3dSamePadding(self):
+    expected_output = [31., 32., 33., 34., 35., 36.]
+    self._VerifyValues(
+        nn_ops.max_pool3d,
+        input_sizes=[1, 2, 2, 3, 3],
+        window=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="SAME",
+        expected=expected_output)
+
+  def testMaxPool3dSamePaddingDifferentStrides(self):
+    expected_output = [2., 5., 8., 18., 21., 24., 34., 37., 40.]
+    self._VerifyValues(
+        nn_ops.max_pool3d,
+        input_sizes=[1, 5, 8, 1, 1],
+        window=[1, 2, 3],
+        strides=[2, 3, 1],
+        padding="SAME",
+        expected=expected_output)
+
+    # Test pooling on a larger input, with different stride and kernel
+    # size for the 'z' dimension.
+
+    # Simulate max pooling in numpy to get the expected output.
+    input_data = np.arange(1, 5 * 27 * 27 * 64 + 1).reshape((5, 27, 27, 64))
+    input_data = np.pad(input_data, [[0, 0], [0, 1], [0, 1], [0, 0]],
+                        mode="constant")
+    expected_output = input_data[:, 1::2, 1::2, :]
+    expected_output[:, -1, :, :] = input_data[:, -2, 1::2, :]
+    expected_output[:, :, -1, :] = input_data[:, 1::2, -2, :]
+    expected_output[:, -1, -1, :] = input_data[:, -2, -2, :]
+
+    self._VerifyValues(
+        nn_ops.max_pool3d,
+        input_sizes=[1, 5, 27, 27, 64],
+        window=[1, 2, 2],
+        strides=[1, 2, 2],
+        padding="SAME",
+        expected=expected_output.flatten())
+
+  def testKernelSmallerThanStride(self):
+    self._VerifyValues(
+        nn_ops.max_pool3d,
+        input_sizes=[1, 3, 3, 3, 1],
+        window=[1, 1, 1],
+        strides=[2, 2, 2],
+        padding="SAME",
+        expected=[1, 3, 7, 9, 19, 21, 25, 27])
+
+    self._VerifyValues(
+        nn_ops.max_pool3d,
+        input_sizes=[1, 7, 7, 7, 1],
+        window=[2, 2, 2],
+        strides=[3, 3, 3],
+        padding="VALID",
+        expected=[58, 61, 79, 82, 205, 208, 226, 229])
+
+    self._VerifyValues(
+        nn_ops.avg_pool3d,
+        input_sizes=[1, 3, 3, 3, 1],
+        window=[1, 1, 1],
+        strides=[2, 2, 2],
+        padding="SAME",
+        expected=[1, 3, 7, 9, 19, 21, 25, 27])
+
+    self._VerifyValues(
+        nn_ops.avg_pool3d,
+        input_sizes=[1, 7, 7, 7, 1],
+        window=[2, 2, 2],
+        strides=[3, 3, 3],
+        padding="VALID",
+        expected=[29.5, 32.5, 50.5, 53.5, 176.5, 179.5, 197.5, 200.5])
+
+  def _VerifyGradient(self,
+                      pool_func,
+                      pool_grad_func,
+                      input_sizes,
+                      ksize,
+                      strides,
+                      padding,
+                      pool_grad_grad_func=None):
+    """Verifies the output values of the pooling gradient function.
+
+    Args:
+      pool_func: Forward pooling function
+      pool_grad_func: Pooling gradient function for pool_grad_func
+      input_sizes: Input tensor dimensions.
+      ksize: The kernel size dimensions
+      strides: The stride dimensions
+      padding: Padding type.
+      pool_grad_grad_func: Second-order gradient function, if available.
+    """
+    ksize = [1] + ksize + [1]
+    strides = [1] + strides + [1]
+    total_size = np.prod(input_sizes)
+    x = np.arange(1, total_size + 1, dtype=np.float32).reshape(input_sizes)
+    with self.session() as sess:
+      # Use the forward pool function to compute some corresponding outputs
+      # (needed for the CPU device, and we need the shape in both cases).
+      with ops.device("CPU"):
+        inputs = array_ops.placeholder(dtypes.float32, shape=input_sizes)
+        outputs = pool_func(
+            inputs,
+            ksize=ksize,
+            strides=strides,
+            padding=padding)
+
+      output_vals = np.array(sess.run(outputs, {inputs: x}))
+      output_gradient_vals = np.arange(
+          1, output_vals.size + 1, dtype=np.float32)
+      output_gradient_vals = output_gradient_vals.reshape(output_vals.shape)
+      output_grad_grad_vals = np.arange(1, x.size + 1, dtype=np.float32)
+      output_grad_grad_vals = output_grad_grad_vals.reshape(x.shape)
+
+      # Use the Tensorflow CPU pooling gradient to compute the expected input
+      # gradients.
+      with ops.device("CPU"):
+        output_gradients = array_ops.placeholder(
+            dtypes.float32, shape=output_vals.shape)
+        expected_input_gradients = pool_grad_func(
+            inputs,
+            outputs,
+            output_gradients,
+            ksize=ksize,
+            strides=strides,
+            padding=padding)
+        expected_input_gradient_vals = sess.run(
+            expected_input_gradients,
+            {inputs: x,
+             output_gradients: output_gradient_vals})
+
+        output_grad_gradients = array_ops.placeholder(
+            dtypes.float32, shape=expected_input_gradient_vals.shape)
+        if pool_grad_grad_func is not None:
+          expected_grad_gradients = pool_grad_grad_func(
+              inputs,
+              outputs,
+              output_grad_gradients,
+              ksize=ksize,
+              strides=strides,
+              padding=padding,
+              data_format="NDHWC")
+          expected_grad_gradients_vals = sess.run(expected_grad_gradients, {
+              inputs: x,
+              output_grad_gradients: output_grad_grad_vals
+          })
+
+      # Run the gradient op on the XLA device
+      with self.test_scope():
+        outputs = array_ops.placeholder(dtypes.float32, shape=output_vals.shape)
+        actual_input_gradients = pool_grad_func(
+            inputs,
+            outputs,
+            output_gradients,
+            ksize=ksize,
+            strides=strides,
+            padding=padding)
+        if pool_grad_grad_func is not None:
+          actual_grad_gradients = pool_grad_grad_func(
+              inputs,
+              outputs,
+              output_grad_gradients,
+              ksize=ksize,
+              strides=strides,
+              padding=padding,
+              data_format="NDHWC")
+
+      actual = sess.run(actual_input_gradients, {
+          inputs: x,
+          outputs: output_vals,
+          output_gradients: output_gradient_vals
+      })
+
+      # Compare the Tensorflow and XLA results.
+      self.assertAllClose(
+          expected_input_gradient_vals.flatten(),
+          actual.flatten(),
+          rtol=1e-5,
+          atol=1e-6)
+      self.assertShapeEqual(actual, inputs)
+
+      if pool_grad_grad_func is not None:
+        actual_grad_gradients_vals = sess.run(
+            actual_grad_gradients, {
+                inputs: x,
+                outputs: output_vals,
+                output_grad_gradients: output_grad_grad_vals
+            })
+
+        # Compare the Tensorflow and XLA results.
+        self.assertAllClose(
+            expected_grad_gradients_vals,
+            actual_grad_gradients_vals,
+            rtol=1e-4,
+            atol=1e-6)
+        self.assertShapeEqual(actual_grad_gradients_vals, outputs)
+
+  def testMaxPoolGradValidPadding1_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[1, 3, 3, 3, 1],
+        ksize=[1, 1, 1],
+        strides=[1, 1, 1],
+        padding="VALID",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testMaxPoolGradValidPadding2_1_6_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[2, 3, 3, 6, 3],
+        ksize=[2, 2, 2],
+        strides=[1, 1, 1],
+        padding="VALID",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testMaxPoolGradValidPadding2_1_7_3d(self):
+    # TODO(b/73062247): the bfloat16 implementation of MaxPool3DGradGrad does
+    # not have enough precision for this test case to pass if
+    # pool_grad_grad_func is passed.
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[2, 3, 5, 7, 3],
+        ksize=[2, 2, 2],
+        strides=[1, 1, 1],
+        padding="VALID")
+
+  def testMaxPoolGradValidPadding2_2_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[2, 2, 2, 2, 3],
+        ksize=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="VALID",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testMaxPoolGradSamePadding1_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[2, 3, 2, 4, 1],
+        ksize=[1, 1, 1],
+        strides=[1, 1, 1],
+        padding="SAME",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testMaxPoolGradSamePadding2_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[2, 3, 2, 4, 1],
+        ksize=[2, 2, 2],
+        strides=[1, 1, 1],
+        padding="SAME",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testMaxPoolGradSamePadding2_2_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[2, 5, 2, 4, 3],
+        ksize=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="SAME",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testMaxPoolGradSamePadding3_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.max_pool3d,
+        gen_nn_ops.max_pool3d_grad,
+        input_sizes=[1, 3, 3, 7, 1],
+        ksize=[3, 3, 3],
+        strides=[1, 1, 1],
+        padding="SAME",
+        pool_grad_grad_func=gen_nn_ops.max_pool3d_grad_grad)
+
+  def testAvgPoolGradValidPadding1_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[2, 3, 3, 3, 3],
+        ksize=[1, 1, 1],
+        strides=[1, 1, 1],
+        padding="VALID")
+
+  def testAvgPoolGradValidPadding2_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[2, 3, 3, 3, 3],
+        ksize=[2, 2, 2],
+        strides=[1, 1, 1],
+        padding="VALID")
+
+  def testAvgPoolGradValidPadding2_2_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[2, 2, 2, 2, 3],
+        ksize=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="VALID")
+
+  def testAvgPoolGradSamePadding1_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[2, 3, 2, 4, 3],
+        ksize=[1, 1, 1],
+        strides=[1, 1, 1],
+        padding="SAME")
+
+  def testAvgPoolGradSamePadding2_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[1, 2, 2, 2, 1],
+        ksize=[2, 2, 2],
+        strides=[1, 1, 1],
+        padding="SAME")
+
+  def testAvgPoolGradSamePadding2_2_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[2, 5, 2, 4, 3],
+        ksize=[2, 2, 2],
+        strides=[2, 2, 2],
+        padding="SAME")
+
+  def testAvgPoolGradSamePadding3_1_3d(self):
+    self._VerifyGradient(
+        nn_ops.avg_pool3d,
+        _AvgPoolGrad,
+        input_sizes=[1, 3, 6, 7, 1],
+        ksize=[3, 3, 3],
+        strides=[1, 1, 1],
+        padding="SAME")
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/pooling_ops_test.py

@@ -0,0 +1,625 @@
+# 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 pooling operations."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.platform import googletest
+
+
+def NHWCToNCHW(input_tensor):
+  """Convert the input from NHWC format to NCHW.
+
+  Args:
+    input_tensor:  a 4-D tensor, or a 4-element array representing the same.
+
+  Returns:
+    the converted tensor or a shape array
+  """
+  if isinstance(input_tensor, tensor.Tensor):
+    return array_ops.transpose(input_tensor, [0, 3, 1, 2])
+  else:
+    return [input_tensor[0], input_tensor[3], input_tensor[1], input_tensor[2]]
+
+
+def NCHWToNHWC(input_tensor):
+  """Convert the input from NCHW format to NHWC.
+
+  Args:
+    input_tensor:  a 4-D tensor, or a 4-element array representing the same.
+
+  Returns:
+    the converted tensor or a shape array
+  """
+  if isinstance(input_tensor, tensor.Tensor):
+    return array_ops.transpose(input_tensor, [0, 2, 3, 1])
+  else:
+    return [input_tensor[0], input_tensor[2], input_tensor[3], input_tensor[1]]
+
+
+def GetTestConfigs():
+  """Get all the valid tests configs to run.
+
+  Returns:
+    all the valid test configs
+  """
+  test_configs = ["NHWC", "NCHW"]
+  return test_configs
+
+
+class PoolingTest(xla_test.XLATestCase):
+
+  def _VerifyOneTest(self, pool_func, input_sizes, ksize, strides, padding,
+                     data_format, expected):
+    """Verifies the output values of the pooling function.
+
+    Args:
+      pool_func: Function to be called, currently only co.MaxPool.
+      input_sizes: Input tensor dimensions.
+      ksize: The kernel size dimensions
+      strides: The stride dimensions
+      padding: Padding type.
+      data_format: The data format we use to run the pooling operation.
+      expected: An array containing the expected operation outputs.
+    """
+    total_size = np.prod(input_sizes)
+    # Initializes the input tensor with array containing incrementing
+    # numbers from 1.
+    x = np.array([f * 1.0 for f in range(1, total_size + 1)], dtype=np.float32)
+    x = x.reshape(input_sizes)
+    with self.session() as sess:
+      with self.test_scope():
+        inputs = array_ops.placeholder(dtypes.float32)
+        t = inputs
+        if data_format == "NCHW":
+          t = NHWCToNCHW(t)
+          ksize = NHWCToNCHW(ksize)
+          strides = NHWCToNCHW(strides)
+        t = pool_func(t,
+                      ksize=ksize,
+                      strides=strides,
+                      padding=padding,
+                      data_format=data_format)
+        if data_format == "NCHW":
+          t = NCHWToNHWC(t)
+      actual = sess.run(t, {inputs: x})
+      self.assertAllClose(expected, actual.flatten(), rtol=1e-5, atol=1e-6)
+
+  def _VerifyValues(self, pool_func, input_sizes, ksize, strides, padding,
+                    expected):
+    """Verifies the output values of the pooling function.
+
+    Args:
+      pool_func: Function to be called, co.MaxPool, co.AvgPool,
+        or the Lua version.
+      input_sizes: Input tensor dimensions.
+      ksize: The kernel size dimensions
+      strides: The stride dimensions
+      padding: Padding type.
+      expected: An array containing the expected operation outputs.
+    """
+    for data_format in GetTestConfigs():
+      self._VerifyOneTest(pool_func, input_sizes, ksize, strides, padding,
+                          data_format, expected)
+
+  def testMaxPoolValidPadding(self):
+    expected_output = [13.0, 14.0, 15.0]
+    self._VerifyValues(nn_ops.max_pool,
+                       input_sizes=[1, 3, 3, 3],
+                       ksize=[1, 2, 2, 1],
+                       strides=[1, 2, 2, 1],
+                       padding="VALID",
+                       expected=expected_output)
+
+  def testMaxPoolSamePadding(self):
+    expected_output = [13.0, 14.0, 15.0, 16.0, 17.0, 18.0]
+    self._VerifyValues(nn_ops.max_pool,
+                       input_sizes=[1, 2, 3, 3],
+                       ksize=[1, 2, 2, 1],
+                       strides=[1, 2, 2, 1],
+                       padding="SAME",
+                       expected=expected_output)
+
+  def testMaxPoolSamePaddingNonSquareWindow(self):
+    # input is:
+    # [1.0, 2.0
+    #  3.0  4.0]
+    #
+    # Window of [x, x] should do:
+    #
+    #  [max(1.0, 2.0), max(2.0, padded0),
+    #   max(3.0, 4.0), max(4.0, padded0)]
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 2, 2, 1],
+        ksize=[1, 1, 2, 1],
+        strides=[1, 1, 1, 1],
+        padding="SAME",
+        expected=[2.0, 2.0, 4.0, 4.0])
+
+  def testMaxPoolValidPaddingUnevenStride(self):
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 4, 4, 1],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 1, 2, 1],
+        padding="VALID",
+        expected=[6.0, 8.0, 10.0, 12.0, 14.0, 16.0])
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 4, 4, 1],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 1, 1],
+        padding="VALID",
+        expected=[6.0, 7.0, 8.0, 14.0, 15.0, 16.0])
+
+  def testMaxPoolSamePaddingFilter4(self):
+    expected_output = [
+        21.0, 22.0, 23.0, 24.0, 29.0, 30.0, 31.0, 32.0, 53.0, 54.0, 55.0, 56.0,
+        61.0, 62.0, 63.0, 64.0
+    ]
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 4, 4, 4],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        expected=expected_output)
+
+  def testMaxPoolSamePaddingFilter8(self):
+    expected_output = [
+        145.0, 146.0, 147.0, 148.0, 149.0, 150.0, 151.0, 152.0, 161.0, 162.0,
+        163.0, 164.0, 165.0, 166.0, 167.0, 168.0, 177.0, 178.0, 179.0, 180.0,
+        181.0, 182.0, 183.0, 184.0, 185.0, 186.0, 187.0, 188.0, 189.0, 190.0,
+        191.0, 192.0, 273.0, 274.0, 275.0, 276.0, 277.0, 278.0, 279.0, 280.0,
+        289.0, 290.0, 291.0, 292.0, 293.0, 294.0, 295.0, 296.0, 305.0, 306.0,
+        307.0, 308.0, 309.0, 310.0, 311.0, 312.0, 313.0, 314.0, 315.0, 316.0,
+        317.0, 318.0, 319.0, 320.0, 401.0, 402.0, 403.0, 404.0, 405.0, 406.0,
+        407.0, 408.0, 417.0, 418.0, 419.0, 420.0, 421.0, 422.0, 423.0, 424.0,
+        433.0, 434.0, 435.0, 436.0, 437.0, 438.0, 439.0, 440.0, 441.0, 442.0,
+        443.0, 444.0, 445.0, 446.0, 447.0, 448.0, 465.0, 466.0, 467.0, 468.0,
+        469.0, 470.0, 471.0, 472.0, 481.0, 482.0, 483.0, 484.0, 485.0, 486.0,
+        487.0, 488.0, 497.0, 498.0, 499.0, 500.0, 501.0, 502.0, 503.0, 504.0,
+        505.0, 506.0, 507.0, 508.0, 509.0, 510.0, 511.0, 512.0
+    ]
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 8, 8, 8],
+        ksize=[1, 3, 3, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        expected=expected_output)
+
+  # Tests for DepthwiseMaxPooling on CPU only.
+  def testDepthwiseMaxPool1x1DepthWindow1(self):
+    # input is:
+    # [1.0, ..., 10.0] along depth,
+    #
+    # We maxpool by depth in patches of 2.
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 1, 1, 10],
+        ksize=[1, 1, 1, 2],
+        strides=[1, 1, 1, 2],
+        padding="SAME",
+        expected=[2.0, 4.0, 6.0, 8.0, 10.0])
+
+  def testDepthwiseMaxPool2x2DepthWindow3(self):
+    # input is:
+    #
+    # a 2x2x6 cube, and we depthwise max across 3 to produce a 2x2x2
+    # output.  Each node has contiguous values, so the depthwise max
+    # should be multiples of 3.0.
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 2, 2, 6],
+        ksize=[1, 1, 1, 3],
+        strides=[1, 1, 1, 3],
+        padding="SAME",
+        expected=[3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0])
+
+  def testKernelSmallerThanStrideValid(self):
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 7, 7, 1],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 3, 3, 1],
+        padding="VALID",
+        expected=[9, 12, 30, 33])
+
+  def testKernelSmallerThanStrideSame(self):
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 3, 3, 1],
+        ksize=[1, 1, 1, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        expected=[1, 3, 7, 9])
+
+    self._VerifyValues(
+        nn_ops.max_pool,
+        input_sizes=[1, 4, 4, 1],
+        ksize=[1, 1, 1, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        expected=[1, 3, 9, 11])
+
+  # Average pooling
+  def testAvgPoolValidPadding(self):
+    expected_output = [7, 8, 9]
+    self._VerifyValues(
+        nn_ops.avg_pool,
+        input_sizes=[1, 3, 3, 3],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 2, 1],
+        padding="VALID",
+        expected=expected_output)
+
+  def testAvgPoolSamePadding(self):
+    expected_output = [7., 8., 9., 11.5, 12.5, 13.5]
+    self._VerifyValues(
+        nn_ops.avg_pool,
+        input_sizes=[1, 2, 3, 3],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        expected=expected_output)
+
+
+class PoolGradTest(xla_test.XLATestCase):
+
+  CPU_DEVICE = "/job:localhost/replica:0/task:0/cpu:0"
+
+  def _VerifyOneTest(self,
+                     pool_func,
+                     pool_grad_func,
+                     input_sizes,
+                     ksize,
+                     strides,
+                     padding,
+                     data_format,
+                     pool_grad_grad_func=None):
+    """Verifies the output values of the pooling gradient function.
+
+    Args:
+      pool_func: Forward pooling function
+      pool_grad_func: Pooling gradient function for pool_grad_func
+      input_sizes: Input tensor dimensions.
+      ksize: The kernel size dimensions
+      strides: The stride dimensions
+      padding: Padding type.
+      data_format: The data format we use to run the pooling operation.
+      pool_grad_grad_func: Second-order gradient function, if available.
+    """
+    total_size = np.prod(input_sizes)
+    # TODO(b/73062247): MaxPoolGradGrad can confuse gradients when x is equally
+    # maximal at 16 bits. Switch to np.random.randn when resolved.
+    x = np.arange(1, total_size + 1, dtype=np.float32)
+    x *= (np.random.randint(2, size=total_size) * 2 - 1)  # Flip signs randomly
+    # Verify some specifically interesting values...
+    x[np.random.choice(total_size)] = np.inf
+    x[np.random.choice(total_size)] = -np.inf
+    # TODO(b/74222344): Fix nan handling for max pool grad.
+    # x[np.random.choice(total_size)] = np.nan
+    x = x.reshape(input_sizes)
+    with self.session() as sess:
+      # Use the forward pool function to compute some corresponding outputs
+      # (needed for the CPU device, and we need the shape in both cases).
+      with ops.device(self.CPU_DEVICE):
+        inputs = array_ops.placeholder(dtypes.float32, shape=input_sizes)
+        outputs = pool_func(
+            inputs,
+            ksize=ksize,
+            strides=strides,
+            padding=padding,
+            data_format="NHWC")
+
+      output_vals = np.array(sess.run(outputs, {inputs: x}))
+      output_gradient_vals = np.arange(
+          1, output_vals.size + 1, dtype=np.float32)
+      output_gradient_vals = output_gradient_vals.reshape(output_vals.shape)
+      output_grad_grad_vals = np.arange(1, x.size + 1, dtype=np.float32)
+      output_grad_grad_vals = output_grad_grad_vals.reshape(x.shape)
+
+      # Use the Tensorflow CPU pooling gradient to compute the expected input
+      # gradients.
+      with ops.device(self.CPU_DEVICE):
+        output_gradients = array_ops.placeholder(
+            dtypes.float32, shape=output_vals.shape)
+        expected_input_gradients = pool_grad_func(
+            inputs,
+            outputs,
+            output_gradients,
+            ksize=ksize,
+            strides=strides,
+            padding=padding,
+            data_format="NHWC")
+        expected_input_gradient_vals = sess.run(
+            expected_input_gradients,
+            {inputs: x,
+             output_gradients: output_gradient_vals})
+
+        output_grad_gradients = array_ops.placeholder(
+            dtypes.float32, shape=expected_input_gradient_vals.shape)
+        if pool_grad_grad_func is not None:
+          expected_grad_gradients = pool_grad_grad_func(
+              inputs,
+              outputs,
+              output_grad_gradients,
+              ksize=ksize,
+              strides=strides,
+              padding=padding,
+              data_format="NHWC")
+          expected_grad_gradients_vals = sess.run(expected_grad_gradients, {
+              inputs: x,
+              output_grad_gradients: output_grad_grad_vals
+          })
+
+      # Run the gradient op on the XLA device
+      with self.test_scope():
+        outputs = array_ops.placeholder(dtypes.float32, shape=output_vals.shape)
+        xla_inputs = inputs
+        xla_outputs = outputs
+        xla_output_gradients = output_gradients
+        xla_output_grad_gradients = output_grad_gradients
+        xla_ksize = ksize
+        xla_strides = strides
+        if data_format == "NCHW":
+          xla_inputs = NHWCToNCHW(inputs)
+          xla_outputs = NHWCToNCHW(outputs)
+          xla_output_gradients = NHWCToNCHW(output_gradients)
+          xla_output_grad_gradients = NHWCToNCHW(output_grad_gradients)
+          xla_ksize = NHWCToNCHW(ksize)
+          xla_strides = NHWCToNCHW(strides)
+        actual_input_gradients = pool_grad_func(
+            xla_inputs,
+            xla_outputs,
+            xla_output_gradients,
+            ksize=xla_ksize,
+            strides=xla_strides,
+            padding=padding,
+            data_format=data_format)
+        if data_format == "NCHW":
+          actual_input_gradients = NCHWToNHWC(actual_input_gradients)
+        if pool_grad_grad_func is not None:
+          actual_grad_gradients = pool_grad_grad_func(
+              xla_inputs,
+              xla_outputs,
+              xla_output_grad_gradients,
+              ksize=xla_ksize,
+              strides=xla_strides,
+              padding=padding,
+              data_format=data_format)
+          if data_format == "NCHW":
+            actual_grad_gradients = NCHWToNHWC(actual_grad_gradients)
+      actual_input_gradients_vals = sess.run(actual_input_gradients, {
+          inputs: x,
+          outputs: output_vals,
+          output_gradients: output_gradient_vals
+      })
+      # Compare the Tensorflow and XLA results.
+      self.assertAllClose(
+          expected_input_gradient_vals,
+          actual_input_gradients_vals,
+          rtol=1e-4,
+          atol=1e-6)
+      self.assertShapeEqual(actual_input_gradients_vals, inputs)
+
+      if pool_grad_grad_func is not None:
+        actual_grad_gradients_vals = sess.run(
+            actual_grad_gradients, {
+                inputs: x,
+                outputs: output_vals,
+                output_grad_gradients: output_grad_grad_vals
+            })
+
+        # Compare the Tensorflow and XLA results.
+        self.assertAllClose(
+            expected_grad_gradients_vals,
+            actual_grad_gradients_vals,
+            rtol=1e-4,
+            atol=1e-6)
+        self.assertShapeEqual(actual_grad_gradients_vals, outputs)
+
+  def _VerifyValues(self,
+                    pool_func,
+                    pool_grad_func,
+                    input_sizes,
+                    ksize,
+                    strides,
+                    padding,
+                    pool_grad_grad_func=None):
+    """Verifies the output values of the pooling function.
+
+    Args:
+      pool_func: Pooling function to be called, e.g., tf.nn.max_pool2d
+      pool_grad_func: Corresponding pooling gradient function.
+      input_sizes: Input tensor dimensions.
+      ksize: The kernel size dimensions
+      strides: The stride dimensions
+      padding: Padding type.
+      pool_grad_grad_func: Second-order gradient function, if available.
+    """
+    for data_format in GetTestConfigs():
+      self._VerifyOneTest(
+          pool_func,
+          pool_grad_func,
+          input_sizes,
+          ksize,
+          strides,
+          padding,
+          data_format,
+          pool_grad_grad_func=pool_grad_grad_func)
+
+  def _TestPooling(self, forward_op, backward_op, pool_grad_grad_func=None):
+    # VALID padding
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 3, 3, 3],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 2, 1],
+        padding="VALID",
+        pool_grad_grad_func=pool_grad_grad_func)
+
+    # SAME padding
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 2, 3, 3],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        pool_grad_grad_func=pool_grad_grad_func)
+
+    # SAME padding, non square window
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 2, 2, 1],
+        ksize=[1, 1, 2, 1],
+        strides=[1, 1, 1, 1],
+        padding="SAME",
+        pool_grad_grad_func=pool_grad_grad_func)
+
+    # VALID padding, uneven stride
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 4, 4, 1],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 1, 2, 1],
+        padding="VALID",
+        pool_grad_grad_func=pool_grad_grad_func)
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 4, 4, 1],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 1, 1],
+        padding="VALID",
+        pool_grad_grad_func=pool_grad_grad_func)
+
+    # SAME padding, size 4 input
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 4, 4, 4],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        pool_grad_grad_func=pool_grad_grad_func)
+
+    # SAME padding, size 8 input
+    self._VerifyValues(
+        forward_op,
+        backward_op,
+        input_sizes=[1, 8, 8, 8],
+        ksize=[1, 3, 3, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME",
+        pool_grad_grad_func=pool_grad_grad_func)
+
+  def testMaxPool(self):
+    self._TestPooling(
+        nn_ops.max_pool,
+        gen_nn_ops.max_pool_grad,
+        pool_grad_grad_func=gen_nn_ops.max_pool_grad_grad)
+
+  def testAvgPool(self):
+    # Wrapper around AvgPoolGrad that ignores extra arguments needed by
+    # MaxPoolGrad.
+    def AvgPoolGrad(inputs, outputs, output_gradients, ksize, strides, padding,
+                    data_format):
+      del outputs  # Unused by average-pooling gradients.
+      return gen_nn_ops.avg_pool_grad(
+          inputs.get_shape().as_list(),
+          output_gradients,
+          ksize=ksize,
+          strides=strides,
+          padding=padding,
+          data_format=data_format)
+
+    self._TestPooling(nn_ops.avg_pool, AvgPoolGrad)
+
+  @test_util.disable_mlir_bridge(
+      "TODO(b/266613412): investigate FPE in AvgPoolGrad for TPU"
+  )
+  def testAvgPoolGradSamePaddingZeroStrideZeroSize(self):
+    output_gradient_vals = np.array([0.39117979], dtype=np.float32)
+    output_gradient_vals = output_gradient_vals.reshape([1, 1, 1, 1])
+    with self.session() as sess:
+      with self.test_scope():
+        output_gradients = array_ops.placeholder(
+            dtypes.float32, shape=output_gradient_vals.shape
+        )
+        t = gen_nn_ops.avg_pool_grad(
+            orig_input_shape=[1, 0, 0, 0],
+            grad=output_gradients,
+            ksize=[1, 0, 0, 0],
+            strides=[1, 0, 0, 0],
+            padding="SAME",
+            data_format="NCHW",
+        )
+      with self.assertRaisesRegex(
+          errors.InvalidArgumentError,
+          (
+              "Sliding window ksize field for dimension 1 must be positive but"
+              " is 0"
+          ),
+      ):
+        sess.run(t, {output_gradients: output_gradient_vals})
+
+  # The CPU implementation of AvgPoolGrad doesn't accept kernels smaller than
+  # the stride size, so we only run the following tests on MaxPoolGrad.
+
+  def testMaxPoolKernelSmallerThanStrideValid(self):
+    self._VerifyValues(
+        nn_ops.max_pool,
+        gen_nn_ops.max_pool_grad,
+        input_sizes=[1, 7, 7, 1],
+        ksize=[1, 2, 2, 1],
+        strides=[1, 3, 3, 1],
+        padding="VALID")
+
+  def testMaxPoolKernelSmallerThanStrideSame(self):
+    self._VerifyValues(
+        nn_ops.max_pool,
+        gen_nn_ops.max_pool_grad,
+        input_sizes=[1, 3, 3, 1],
+        ksize=[1, 1, 1, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME")
+
+    self._VerifyValues(
+        nn_ops.max_pool,
+        gen_nn_ops.max_pool_grad,
+        input_sizes=[1, 4, 4, 1],
+        ksize=[1, 1, 1, 1],
+        strides=[1, 2, 2, 1],
+        padding="SAME")
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/proximal_adagrad_test.py

@@ -0,0 +1,170 @@
+# Copyright 2015 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 Proximal Adagrad optimizer."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import adagrad
+from tensorflow.python.training import proximal_adagrad
+
+
+class ProximalAdagradOptimizerTest(xla_test.XLATestCase):
+
+  def testResourceProximalAdagradwithoutRegularization(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([0.0, 0.0])
+      var1 = resource_variable_ops.ResourceVariable([0.0, 0.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.0,
+          l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+      self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+      # Run 3 steps Proximal Adagrad.
+      for _ in range(3):
+        update.run()
+
+      self.assertAllClose(
+          np.array([-2.60260963, -4.29698515]), self.evaluate(var0))
+      self.assertAllClose(
+          np.array([-0.28432083, -0.56694895]), self.evaluate(var1))
+      opt_vars = opt.variables()
+      self.assertStartsWith(opt_vars[0].name, var0._shared_name)
+      self.assertStartsWith(opt_vars[1].name, var1._shared_name)
+      self.assertEqual(2, len(opt_vars))
+
+  def testProximalAdagradwithoutRegularization2(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+      var1 = resource_variable_ops.ResourceVariable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.0,
+          l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+      # Run 3 steps Proximal Adagrad.
+      for _ in range(3):
+        update.run()
+      self.assertAllClose(np.array([-1.60261, -2.296985]), self.evaluate(var0))
+      self.assertAllClose(np.array([3.715679, 2.433051]), self.evaluate(var1))
+
+  def testProximalAdagradWithL1(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+      var1 = resource_variable_ops.ResourceVariable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.001,
+          l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+      # Run 10 steps Proximal Adagrad
+      for _ in range(10):
+        update.run()
+      self.assertAllClose(np.array([-6.663634, -9.190331]), self.evaluate(var0))
+      self.assertAllClose(np.array([2.959304, 1.029232]), self.evaluate(var1))
+
+  def testProximalAdagradWithL1_L2(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+      var1 = resource_variable_ops.ResourceVariable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_adagrad.ProximalAdagradOptimizer(
+          3.0,
+          initial_accumulator_value=0.1,
+          l1_regularization_strength=0.001,
+          l2_regularization_strength=2.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+      # Run 10 steps Proximal Adagrad.
+      for _ in range(10):
+        update.run()
+
+      self.assertAllClose(np.array([-0.0495, -0.0995]), self.evaluate(var0))
+      self.assertAllClose(np.array([-0.0045, -0.0095]), self.evaluate(var1))
+
+  def applyOptimizer(self, opt, steps=5):
+    var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+    var1 = resource_variable_ops.ResourceVariable([3.0, 4.0])
+    grads0 = constant_op.constant([0.1, 0.2])
+    grads1 = constant_op.constant([0.01, 0.02])
+
+    update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+
+    self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+    self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+    # Run ProximalAdagrad for a few steps
+    for _ in range(steps):
+      update.run()
+
+    return self.evaluate(var0), self.evaluate(var1)
+
+  def testEquivAdagradwithoutRegularization(self):
+    with self.session(), self.test_scope():
+      val0, val1 = self.applyOptimizer(
+          proximal_adagrad.ProximalAdagradOptimizer(
+              3.0,
+              initial_accumulator_value=0.1,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0))
+
+    with self.session(), self.test_scope():
+      val2, val3 = self.applyOptimizer(
+          adagrad.AdagradOptimizer(
+              3.0, initial_accumulator_value=0.1))
+
+    self.assertAllClose(val0, val2)
+    self.assertAllClose(val1, val3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/proximal_gradient_descent_test.py

@@ -0,0 +1,152 @@
+# Copyright 2015 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 Proximal Gradient Descent optimizer."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import gradient_descent
+from tensorflow.python.training import proximal_gradient_descent
+
+
+class ProximalGradientDescentOptimizerTest(xla_test.XLATestCase):
+
+  def testResourceProximalGradientDescentwithoutRegularization(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([0.0, 0.0])
+      var1 = resource_variable_ops.ResourceVariable([0.0, 0.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.0, l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([0.0, 0.0], self.evaluate(var0))
+      self.assertAllClose([0.0, 0.0], self.evaluate(var1))
+
+      # Run 3 steps Proximal Gradient Descent.
+      for _ in range(3):
+        update.run()
+
+      self.assertAllClose(np.array([-0.9, -1.8]), self.evaluate(var0))
+      self.assertAllClose(np.array([-0.09, -0.18]), self.evaluate(var1))
+
+  def testProximalGradientDescentwithoutRegularization2(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+      var1 = resource_variable_ops.ResourceVariable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.0, l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+      # Run 3 steps Proximal Gradient Descent
+      for _ in range(3):
+        update.run()
+
+      self.assertAllClose(np.array([0.1, 0.2]), self.evaluate(var0))
+      self.assertAllClose(np.array([3.91, 2.82]), self.evaluate(var1))
+
+  def testProximalGradientDescentWithL1(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+      var1 = resource_variable_ops.ResourceVariable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.001, l2_regularization_strength=0.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+      # Run 10 steps proximal gradient descent.
+      for _ in range(10):
+        update.run()
+
+      self.assertAllClose(np.array([-1.988, -3.988001]), self.evaluate(var0))
+      self.assertAllClose(np.array([3.67, 2.37]), self.evaluate(var1))
+
+  def testProximalGradientDescentWithL1_L2(self):
+    with self.session(), self.test_scope():
+      var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+      var1 = resource_variable_ops.ResourceVariable([4.0, 3.0])
+      grads0 = constant_op.constant([0.1, 0.2])
+      grads1 = constant_op.constant([0.01, 0.02])
+
+      opt = proximal_gradient_descent.ProximalGradientDescentOptimizer(
+          3.0, l1_regularization_strength=0.001, l2_regularization_strength=2.0)
+      update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+      self.evaluate(variables.global_variables_initializer())
+
+      self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+      self.assertAllClose([4.0, 3.0], self.evaluate(var1))
+
+      # Run 10 steps Proximal Gradient Descent
+      for _ in range(10):
+        update.run()
+
+      self.assertAllClose(np.array([-0.0495, -0.0995]), self.evaluate(var0))
+      self.assertAllClose(np.array([-0.0045, -0.0095]), self.evaluate(var1))
+
+  def applyOptimizer(self, opt, steps=5):
+    var0 = resource_variable_ops.ResourceVariable([1.0, 2.0])
+    var1 = resource_variable_ops.ResourceVariable([3.0, 4.0])
+    grads0 = constant_op.constant([0.1, 0.2])
+    grads1 = constant_op.constant([0.01, 0.02])
+
+    update = opt.apply_gradients(zip([grads0, grads1], [var0, var1]))
+    self.evaluate(variables.global_variables_initializer())
+
+    self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+    self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+    # Run ProximalAdagrad for a few steps
+    for _ in range(steps):
+      update.run()
+
+    return self.evaluate(var0), self.evaluate(var1)
+
+  def testEquivGradientDescentwithoutRegularization(self):
+    with self.session(), self.test_scope():
+      val0, val1 = self.applyOptimizer(
+          proximal_gradient_descent.ProximalGradientDescentOptimizer(
+              3.0,
+              l1_regularization_strength=0.0,
+              l2_regularization_strength=0.0))
+
+    with self.session(), self.test_scope():
+      val2, val3 = self.applyOptimizer(
+          gradient_descent.GradientDescentOptimizer(3.0))
+
+    self.assertAllClose(val0, val2)
+    self.assertAllClose(val1, val3)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/qr_op_test.py

@@ -0,0 +1,148 @@
+# Copyright 2018 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 tensorflow.ops.math_ops.matrix_inverse."""
+
+import itertools
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+@test_util.run_all_without_tensor_float_32(
+    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
+    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
+    "occasionally occurs")
+# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
+class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def AdjustedNorm(self, x):
+    """Computes the norm of matrices in 'x', adjusted for dimension and type."""
+    norm = np.linalg.norm(x, axis=(-2, -1))
+    return norm / (max(x.shape[-2:]) * np.finfo(x.dtype).eps)
+
+  def CompareOrthogonal(self, x, y, rank):
+    # We only compare the first 'rank' orthogonal vectors since the
+    # remainder form an arbitrary orthonormal basis for the
+    # (row- or column-) null space, whose exact value depends on
+    # implementation details. Notice that since we check that the
+    # matrices of singular vectors are unitary elsewhere, we do
+    # implicitly test that the trailing vectors of x and y span the
+    # same space.
+    x = x[..., 0:rank]
+    y = y[..., 0:rank]
+    # Q is only unique up to sign (complex phase factor for complex matrices),
+    # so we normalize the sign first.
+    sum_of_ratios = np.sum(np.divide(y, x), -2, keepdims=True)
+    phases = np.divide(sum_of_ratios, np.abs(sum_of_ratios))
+    x *= phases
+    self.assertTrue(np.all(self.AdjustedNorm(x - y) < 30.0))
+
+  def CheckApproximation(self, a, q, r):
+    # Tests that a ~= q*r.
+    precision = self.AdjustedNorm(a - np.matmul(q, r))
+    self.assertTrue(np.all(precision < 11.0))
+
+  def CheckUnitary(self, x):
+    # Tests that x[...,:,:]^H * x[...,:,:] is close to the identity.
+    xx = math_ops.matmul(x, x, adjoint_a=True)
+    identity = array_ops.matrix_band_part(array_ops.ones_like(xx), 0, 0)
+    tol = 100 * np.finfo(x.dtype).eps
+    self.assertAllClose(xx, identity, atol=tol)
+
+  def _random_matrix(self, dtype, shape):
+    np.random.seed(1)
+
+    def rng():
+      return np.random.uniform(
+          low=-1.0, high=1.0, size=np.prod(shape)).reshape(shape).astype(dtype)
+
+    x_np = rng()
+    if np.issubdtype(dtype, np.complexfloating):
+      x_np += rng() * dtype(1j)
+    return x_np
+
+  def _test(self, x_np, full_matrices, full_rank=True):
+    dtype = x_np.dtype
+    shape = x_np.shape
+    with self.session() as sess:
+      x_tf = array_ops.placeholder(dtype)
+      with self.device_scope():
+        q_tf, r_tf = linalg_ops.qr(x_tf, full_matrices=full_matrices)
+      q_tf_val, r_tf_val = sess.run([q_tf, r_tf], feed_dict={x_tf: x_np})
+
+      q_dims = q_tf_val.shape
+      np_q = np.ndarray(q_dims, dtype)
+      np_q_reshape = np.reshape(np_q, (-1, q_dims[-2], q_dims[-1]))
+      new_first_dim = np_q_reshape.shape[0]
+
+      x_reshape = np.reshape(x_np, (-1, x_np.shape[-2], x_np.shape[-1]))
+      for i in range(new_first_dim):
+        if full_matrices:
+          np_q_reshape[i, :, :], _ = np.linalg.qr(
+              x_reshape[i, :, :], mode="complete")
+        else:
+          np_q_reshape[i, :, :], _ = np.linalg.qr(
+              x_reshape[i, :, :], mode="reduced")
+      np_q = np.reshape(np_q_reshape, q_dims)
+      if full_rank:
+        # Q is unique up to sign/phase if the matrix is full-rank.
+        self.CompareOrthogonal(np_q, q_tf_val, min(shape[-2:]))
+      self.CheckApproximation(x_np, q_tf_val, r_tf_val)
+      self.CheckUnitary(q_tf_val)
+
+  SIZES = [1, 2, 5, 10, 32, 100, 300, 603]
+  DTYPES = [np.float32, np.complex64]
+  PARAMS = itertools.product(SIZES, SIZES, DTYPES)
+
+  @parameterized.parameters(*PARAMS)
+  def testQR(self, rows, cols, dtype):
+    for full_matrices in [True, False]:
+      # Only tests the (3, 2) case for small numbers of rows/columns.
+      for batch_dims in [(), (3,)] + [(3, 2)] * (max(rows, cols) < 10):
+        x_np = self._random_matrix(dtype, batch_dims + (rows, cols))
+        self._test(x_np, full_matrices)
+
+  def testLarge2000x2000(self):
+    x_np = self._random_matrix(np.float32, (2000, 2000))
+    self._test(x_np, full_matrices=True)
+
+  @unittest.skip("Test times out on CI")
+  def testLarge17500x128(self):
+    x_np = self._random_matrix(np.float32, (17500, 128))
+    self._test(x_np, full_matrices=True)
+
+  @parameterized.parameters((23, 25), (513, 23))
+  def testZeroColumn(self, rows, cols):
+    x_np = self._random_matrix(np.complex64, (rows, cols))
+    x_np[:, 7] = 0.
+    self._test(x_np, full_matrices=True)
+
+  @parameterized.parameters((4, 4), (514, 20))
+  def testRepeatedColumn(self, rows, cols):
+    x_np = self._random_matrix(np.complex64, (rows, cols))
+    x_np[:, 1] = x_np[:, 2]
+    self._test(x_np, full_matrices=True, full_rank=False)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/quantized_ops_test.py

@@ -0,0 +1,100 @@
+# Copyright 2018 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 quantized operations."""
+
+import math
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.compiler.tf2xla.python import xla
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import bitwise_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class QuantizedOpsTest(xla_test.XLATestCase):
+
+  # Verify that quantized types can be clustered by XLA.
+  def testQuantizedTypeRoundtrip(self):
+    with self.session() as session:
+      for dtype in self.quantized_tf_types:
+        in_values = np.array([1, 2, 3, 4, 5, 6])
+        expected = [[1, 2], [3, 4], [5, 6]]
+        with self.test_scope():
+          p = array_ops.placeholder(dtype=dtypes.int32)
+          x = math_ops.cast(p, dtype)
+          x = array_ops.reshape(x, [3, 2])
+
+        value = session.run(x, {p: in_values})
+        self.assertAllEqual(value, expected)
+
+
+class DequantizedOpsTest(xla_test.XLATestCase):
+
+  def pack_uint8_r2_to_uint32(self, test_input):
+    num_rows, num_columns = test_input.get_shape().as_list()
+    num_output_columns = int(math.ceil(num_columns / 4.0))
+    padding_input = array_ops.pad(
+        math_ops.cast(test_input, dtype=dtypes.uint8),
+        constant_op.constant([[
+            0,
+            0,
+        ], [0, num_output_columns * 4 - num_columns]]))
+    output = array_ops.zeros([num_rows, num_output_columns],
+                             dtype=dtypes.uint32)
+    num_elements_per_pack = 4
+    shift_bits = 8
+
+    iota_r1 = math_ops.range(num_output_columns * num_elements_per_pack)
+
+    for p in range(num_elements_per_pack):
+      selected_index = math_ops.equal(
+          math_ops.mod(iota_r1, num_elements_per_pack), p)
+      gather_index = array_ops.boolean_mask(iota_r1, selected_index)
+      gathered_input = array_ops.gather(padding_input, gather_index, axis=1)
+      total_shift_bits = shift_bits * (num_elements_per_pack - p - 1)
+      left_shift_input = bitwise_ops.left_shift(
+          math_ops.cast(gathered_input, dtype=dtypes.uint32), total_shift_bits)
+      output = bitwise_ops.bitwise_or(output, left_shift_input)
+    return output
+
+  def testDequantizeQuint8(self):
+    num_rows = 100
+    num_columns = 3547
+    random_input = np.random.normal(128.0, 10.0, [num_rows, num_columns])
+    with self.session() as session:
+      with ops.device("CPU"):
+        test_input = ops.convert_to_tensor(random_input, dtype=dtypes.float32)
+        transposed_input = array_ops.transpose(test_input, [1, 0])
+        quantized_input = array_ops.quantize(transposed_input, 0.0, 255.0,
+                                             dtypes.quint8)
+        packed_input = self.pack_uint8_r2_to_uint32(quantized_input.output)
+      with self.test_scope():
+        transposed_quantized_output = xla.dequantize(packed_input, 0.0, 255.0,
+                                                     "MIN_COMBINED", True)
+        quantized_output = array_ops.slice(transposed_quantized_output, [0, 0],
+                                           [num_rows, num_columns])
+
+      value = session.run(quantized_output)
+    self.assertAllClose(value, random_input, 1.0)
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/random_ops_test.py

@@ -0,0 +1,290 @@
+# 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.
+# ==============================================================================
+"""Tests for random-number generation ops in the XLA JIT compiler."""
+
+import math
+
+from absl.testing import parameterized
+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.ops import random_ops
+from tensorflow.python.ops.distributions import special_math
+from tensorflow.python.platform import googletest
+
+
+class RandomOpsTest(xla_test.XLATestCase, parameterized.TestCase):
+  """Test cases for random-number generating operators."""
+
+  def _random_types(self):
+    return set(self.numeric_types) - set(
+        self.complex_types) - {np.uint64, np.int64, np.uint8, np.int8}
+
+  def _testRngIsNotConstant(self, rng, dtype):
+    # Tests that 'rng' does not always return the same value.
+    with self.session():
+      with self.test_scope():
+        x = rng(dtype)
+
+      # The random-number generator, if working correctly, should produce the
+      # same output multiple times with low probability.
+      y = self.evaluate(x)
+      z = self.evaluate(x)
+      w = self.evaluate(x)
+
+      # We use exact equality here. If the random-number generator is producing
+      # deterministic output, all three outputs will be bitwise identical.
+      self.assertTrue((not np.array_equal(y, z)) or
+                      (not np.array_equal(z, w)) or (not np.array_equal(y, w)))
+
+  def testRandomUniformIsNotConstant(self):
+
+    def rng(dtype):
+      dtype = dtypes.as_dtype(dtype)
+      return random_ops.random_uniform(shape=[2], dtype=dtype, maxval=dtype.max)
+
+    for dtype in self._random_types():
+      self._testRngIsNotConstant(rng, dtype)
+
+  def testRandomNormalIsNotConstant(self):
+
+    def rng(dtype):
+      return random_ops.random_normal(shape=[2], dtype=dtype)
+
+    for dtype in self._random_types() & self.float_types:
+      self._testRngIsNotConstant(rng, dtype)
+
+  @parameterized.parameters({
+      'mean': 1.4,
+      'stddev': 1.2
+  }, {
+      'mean': 2.3,
+      'stddev': 2.0
+  })
+  def testRandomNormal(self, mean, stddev):
+    num_elts = 1000000
+    for dtype in self._random_types() & self.float_types:
+      with self.session():
+        with self.test_scope():
+          normal = random_ops.random_normal([num_elts],
+                                            dtype=dtype,
+                                            mean=mean,
+                                            stddev=stddev)
+          self._checkTruncatedNormalIsInRange(
+              normal,
+              a=normal.dtype.min,
+              b=normal.dtype.max,
+              mu=mean,
+              sigma=stddev,
+              count=num_elts,
+              stat_test=True)
+
+  def testRandomUniformIsInRange(self):
+    for dtype in self._random_types():
+      # TODO (b/112272078): enable bfloat16 for CPU and GPU when the bug is
+      # fixed.
+      if (self.device in ['XLA_GPU', 'XLA_CPU'
+                         ]) and (dtype in [dtypes.bfloat16, dtypes.half]):
+        continue
+      with self.session():
+        with self.test_scope():
+          x = random_ops.random_uniform(
+              shape=[1000], dtype=dtype, minval=-2, maxval=33)
+        y = self.evaluate(x)
+        msg = str(y) + str(dtype)
+        self.assertEqual((y >= -2).sum(), 1000, msg)
+        self.assertEqual((y < 33).sum(), 1000, msg)
+
+  def testTruncatedNormalIsNotConstant(self):
+
+    def rng(dtype):
+      return random_ops.truncated_normal(shape=[2], dtype=dtype)
+
+    # TODO(b/34339814): make this test work with 16 bit float types.
+    for dtype in self._random_types() & {np.float32, np.float64}:
+      self._testRngIsNotConstant(rng, dtype)
+
+  def _checkTruncatedNormalIsInRange(self, x, a, b, mu, sigma, count,
+                                     stat_test):
+
+    def normal_cdf(x):
+      return .5 * math.erfc(-x / math.sqrt(2))
+
+    def normal_pdf(x):
+      return math.exp(-(x**2) / 2.) / math.sqrt(2 * math.pi)
+
+    def probit(x):
+      return self.evaluate(special_math.ndtri(x))
+
+    y = self.evaluate(x)
+
+    alpha = (a - mu) / sigma
+    beta = (b - mu) / sigma
+    z = normal_cdf(beta) - normal_cdf(alpha)
+
+    self.assertEqual((y >= a).sum(), count)
+    self.assertEqual((y <= b).sum(), count)
+
+    # Skip statistical test for low probability regions.
+    if not stat_test:
+      return
+
+    # For more information on these calculations, see:
+    # Burkardt, John. "The Truncated Normal Distribution".
+    # Department of Scientific Computing website. Florida State University.
+    expected_mean = mu + (normal_pdf(alpha) - normal_pdf(beta)) / z * sigma
+    actual_mean = np.mean(y, dtype=np.float64)
+    if x.dtype == dtypes.bfloat16:
+      atol = rtol = 1e-1
+    else:
+      atol = rtol = 2e-2
+    self.assertAllClose(actual_mean, expected_mean, atol=atol, rtol=rtol)
+
+    expected_median = mu + probit(
+        (normal_cdf(alpha) + normal_cdf(beta)) / 2.) * sigma
+    actual_median = np.median(y)
+    self.assertAllClose(actual_median, expected_median, atol=atol, rtol=rtol)
+
+    expected_variance = sigma**2 * (1 + (
+        (alpha * normal_pdf(alpha) - beta * normal_pdf(beta)) / z) - (
+            (normal_pdf(alpha) - normal_pdf(beta)) / z)**2)
+    actual_variance = np.var(y, dtype=np.float64)
+    self.assertAllClose(
+        actual_variance, expected_variance, atol=atol, rtol=rtol)
+
+  def testTruncatedNormalIsInRange(self):
+    count = 10000000
+    # TODO(b/34339814): make this test work with 16 bit float types.
+    for dtype in self._random_types() & {np.float32, np.float64}:
+      with self.session():
+        with self.test_scope():
+          x = random_ops.truncated_normal(shape=[count], dtype=dtype)
+        self._checkTruncatedNormalIsInRange(
+            x, a=-2, b=2, mu=0, sigma=1, count=count, stat_test=True)
+
+  def _implParameterizedTruncatedNormalIsInRange(self, a, b, mu, sigma, count,
+                                                 stat_test):
+    # TODO(b/34339814): make this test work with 16 bit float types.
+    for dtype in self._random_types() & {np.float32, np.float64}:
+      with self.session():
+        with self.test_scope():
+          x = random_ops.parameterized_truncated_normal(
+              shape=[count],
+              dtype=dtype,
+              means=mu,
+              stddevs=sigma,
+              minvals=a,
+              maxvals=b)
+        self._checkTruncatedNormalIsInRange(
+            x, a=a, b=b, mu=mu, sigma=sigma, count=count, stat_test=stat_test)
+
+  def testParameterizedTruncatedNormalBroadcasting(self):
+    for dtype in self._random_types() & {np.float32, np.float64}:
+      with self.session():
+        with self.test_scope():
+          a = -1.
+          b = 1.
+          mu = 0.
+          sigma = 1.
+          count = 10000000
+          x = random_ops.parameterized_truncated_normal(
+              shape=[1, count],
+              dtype=dtype,
+              means=mu,
+              stddevs=sigma,
+              minvals=[a],
+              maxvals=[b])
+        self._checkTruncatedNormalIsInRange(
+            x, a=a, b=b, mu=mu, sigma=sigma, count=count, stat_test=True)
+
+  def testParameterizedTruncatedNormalBatched(self):
+    # TODO(b/112289993): Make this test work with dtype np.float64.
+    for dtype in self._random_types() & {np.float32}:
+      with self.session():
+        with self.test_scope():
+          count = 10000000
+          a = -100.
+          b = 100.
+          mu0 = 0.
+          mu1 = 1.
+          sigma = .1
+          x = random_ops.parameterized_truncated_normal(
+              shape=[2, count],
+              dtype=dtype,
+              means=[mu0, mu1],
+              stddevs=sigma,
+              minvals=[a],
+              maxvals=[b])
+        self._checkTruncatedNormalIsInRange(
+            x[0], a=a, b=b, mu=mu0, sigma=sigma, count=count, stat_test=True)
+        self._checkTruncatedNormalIsInRange(
+            x[1], a=a, b=b, mu=mu1, sigma=sigma, count=count, stat_test=True)
+
+  def testParameterizedTruncatedNormalIsInRangeCenter(self):
+    count = 10000000
+    self._implParameterizedTruncatedNormalIsInRange(
+        a=-10, b=20, mu=5, sigma=5, count=count, stat_test=True)
+
+  def testParameterizedTruncatedNormalIsInRangeLeft(self):
+    count = 10000000
+    # the region is on the left side of the parent normal distribution
+    self._implParameterizedTruncatedNormalIsInRange(
+        a=-10, b=-4, mu=0, sigma=1, count=count, stat_test=False)
+    self._implParameterizedTruncatedNormalIsInRange(
+        a=-np.inf, b=-4, mu=0, sigma=1, count=count, stat_test=False)
+
+  def testParameterizedTruncatedNormalIsInRangeRight(self):
+    count = 10000000
+    # the region is on the right side of the parent normal distribution
+    self._implParameterizedTruncatedNormalIsInRange(
+        a=4, b=10, mu=0, sigma=1, count=count, stat_test=False)
+    self._implParameterizedTruncatedNormalIsInRange(
+        a=4, b=np.inf, mu=0, sigma=1, count=count, stat_test=False)
+
+  def testShuffle1d(self):
+    with self.session():
+      with self.test_scope():
+        x = math_ops.range(1 << 16)
+        shuffle = random_ops.random_shuffle(x)
+      result = self.evaluate(shuffle)
+      expected = range(1 << 16)
+      # Compare sets to avoid randomness behavior changes but make sure still
+      # have all the values.
+      self.assertAllEqual(set(result), set(expected))
+
+  def testShuffle2d(self):
+    with self.session():
+      with self.test_scope():
+        x = array_ops.diag(math_ops.range(20))
+        shuffle = random_ops.random_shuffle(x)
+      result = self.evaluate(shuffle)
+      expected = np.diag(range(20)).flatten()
+      # Compare sets to avoid randomness behavior changes but make sure still
+      # have all the values.
+      self.assertAllEqual(len(result.flatten()), len(expected))
+      self.assertAllEqual(set(result.flatten()), set(expected))
+
+  def testRandomShuffleInputRank0(self):
+    with self.session():
+      with self.test_scope():
+        shuffle = random_ops.random_shuffle(value=1e20)
+      self.evaluate(shuffle)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/reduce_ops_test.py

@@ -0,0 +1,248 @@
+# 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.
+# ==============================================================================
+"""Tests for reduction operators."""
+
+import functools
+import itertools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+@parameterized.named_parameters(('32_bit_index', dtypes.int32),
+                                ('64_bit_index', dtypes.int64))
+class ReduceOpsTest(xla_test.XLATestCase, parameterized.TestCase):
+  def _testReduction(self,
+                     tf_reduce_fn,
+                     np_reduce_fn,
+                     dtype,
+                     test_inputs,
+                     index_dtype,
+                     rtol=1e-4,
+                     atol=1e-4):
+    """Tests that the output of 'tf_reduce_fn' matches numpy's output."""
+
+    for test_input in test_inputs:
+      with self.session() as sess:
+        with self.test_scope():
+          a = array_ops.placeholder(dtype)
+          index = array_ops.placeholder(index_dtype)
+          out = tf_reduce_fn(a, index)
+        result = sess.run(out, {a: test_input, index: [0]})
+        self.assertAllClose(
+            result, np_reduce_fn(test_input, axis=0), rtol=rtol, atol=atol)
+
+        result = sess.run(out, {a: test_input, index: [1]})
+        self.assertAllClose(
+            result, np_reduce_fn(test_input, axis=1), rtol=rtol, atol=atol)
+
+        result = sess.run(out, {a: test_input, index: [-1]})
+        self.assertAllClose(
+            result, np_reduce_fn(test_input, axis=1), rtol=rtol, atol=atol)
+
+        # MLIR bridge doesn't return the same error so it can't be matched
+        # directly.
+        if not test_util.is_mlir_bridge_enabled():
+          with self.assertRaisesWithPredicateMatch(
+              errors_impl.InvalidArgumentError, 'Invalid reduction dim'):
+            sess.run(out, {a: test_input, index: [-33]})
+
+          with self.assertRaisesWithPredicateMatch(
+              errors_impl.InvalidArgumentError, 'Invalid reduction dim'):
+            sess.run(out, {a: test_input, index: [2]})
+
+  REAL_DATA = [
+      np.zeros(shape=(2, 0)),
+      np.zeros(shape=(0, 30)),
+      np.arange(1, 7).reshape(2, 3),
+      np.arange(-10, -4).reshape(2, 3),
+      np.arange(-4, 2).reshape(2, 3),
+  ]
+  COMPLEX_DATA = [
+      np.zeros(shape=(2, 0)).astype(np.complex64),
+      np.zeros(shape=(0, 30)).astype(np.complex64),
+      np.arange(1, 13, dtype=np.float32).view(np.complex64).reshape(2, 3),
+      np.arange(-14, -2, dtype=np.float32).view(np.complex64).reshape(2, 3),
+      np.arange(-4, 8, dtype=np.float32).view(np.complex64).reshape(2, 3),
+  ]
+  NONEMPTY_REAL_DATA = [x for x in REAL_DATA if np.size(x) > 0]
+  NONEMPTY_COMPLEX_DATA = [x for x in COMPLEX_DATA if np.size(x) > 0]
+  BOOL_DATA = [
+      np.array([], dtype=np.bool_).reshape(2, 0),
+      np.array([], dtype=np.bool_).reshape(0, 3),
+      np.array([[False, True, False], [True, True, False]]),
+  ]
+  ONES = [np.ones([34000, 2])]
+
+  def testReduceSumF32(self, index_dtype):
+    self._testReduction(math_ops.reduce_sum, np.sum, np.float32, self.REAL_DATA,
+                        index_dtype)
+
+  def testReduceSumC64(self, index_dtype):
+    self._testReduction(math_ops.reduce_sum, np.sum, np.complex64,
+                        self.COMPLEX_DATA, index_dtype)
+
+  def testReduceProdF32(self, index_dtype):
+    self._testReduction(math_ops.reduce_prod, np.prod, np.float32,
+                        self.REAL_DATA, index_dtype)
+
+  def testReduceProdC64(self, index_dtype):
+    self._testReduction(math_ops.reduce_prod, np.prod, np.complex64,
+                        self.COMPLEX_DATA, index_dtype)
+
+  def testReduceMin(self, index_dtype):
+
+    def reference_min(dtype, inp, axis):
+      """Wrapper around np.amin that returns +infinity for an empty input."""
+      if inp.shape[axis] == 0:
+        if np.issubdtype(dtype, np.floating):
+          return np.full(inp.shape[0:axis] + inp.shape[axis + 1:], float('inf'))
+        return np.full(inp.shape[0:axis] + inp.shape[axis + 1:],
+                       np.iinfo(dtype).max)
+      return np.amin(inp, axis)
+
+    for dtype in set(self.all_types).intersection(
+        [np.float32, np.int32, np.int64]):
+      self._testReduction(math_ops.reduce_min,
+                          functools.partial(reference_min, dtype), dtype,
+                          self.REAL_DATA, index_dtype)
+
+  def testReduceMax(self, index_dtype):
+
+    def reference_max(dtype, inp, axis):
+      """Wrapper around np.amax that returns -infinity for an empty input."""
+      if inp.shape[axis] == 0:
+        if np.issubdtype(dtype, np.floating):
+          return np.full(inp.shape[0:axis] + inp.shape[axis + 1:],
+                         float('-inf'))
+        return np.full(inp.shape[0:axis] + inp.shape[axis + 1:],
+                       np.iinfo(dtype).min)
+      return np.amax(inp, axis)
+
+    for dtype in set(self.all_types).intersection(
+        [np.float32, np.int32, np.int64]):
+      self._testReduction(math_ops.reduce_max,
+                          functools.partial(reference_max, dtype), dtype,
+                          self.REAL_DATA, index_dtype)
+
+  def testReduceMeanF32(self, index_dtype):
+    # TODO(phawkins): mean on XLA currently returns 0 instead of NaN when
+    # reducing across zero inputs.
+    self._testReduction(math_ops.reduce_mean, np.mean, np.float32,
+                        self.NONEMPTY_REAL_DATA, index_dtype)
+
+  def testReduceMeanF16(self, index_dtype):
+    if np.float16 in self.all_types:
+      self._testReduction(math_ops.reduce_mean, np.mean, np.float16, self.ONES,
+                          index_dtype)
+
+  def testReduceMeanC64(self, index_dtype):
+    self._testReduction(math_ops.reduce_mean, np.mean, np.complex64,
+                        self.NONEMPTY_COMPLEX_DATA, index_dtype)
+
+  def testReduceAll(self, index_dtype):
+    self._testReduction(math_ops.reduce_all, np.all, np.bool_, self.BOOL_DATA,
+                        index_dtype)
+
+  def testReduceAny(self, index_dtype):
+    self._testReduction(math_ops.reduce_any, np.any, np.bool_, self.BOOL_DATA,
+                        index_dtype)
+
+  @test_util.disable_mlir_bridge('Error messages differ')
+  def testReduceSumWithDuplicateAxes(self, index_dtype):
+    with self.session() as sess:
+      with self.test_scope():
+        a = array_ops.placeholder(np.float32)
+        index = array_ops.placeholder(np.int32)
+        out = math_ops.reduce_sum(a, index)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          'Axes contains duplicate dimension'):
+        sess.run(out, {a: [10, 20, 30], index: [0, 0]})
+
+
+class ReduceOpPrecisionTest(xla_test.XLATestCase):
+
+  def _testReduceSum(self,
+                     expected_result,
+                     dtype,
+                     test_inputs,
+                     rtol=1e-3,
+                     atol=1e-4):
+    """Tests reduce sum on a list of input arrays.
+
+    For each array in test_inputs, check that performing reduce sum on the array
+    produces a value that is close to the expected result.
+
+    Args:
+      expected_result: the expected result.
+      dtype: the data type of the reduce sum operation.
+      test_inputs: a list of input arrays for the reduce sum operation.
+      rtol: the relative error.
+      atol: the absolute error.
+    """
+
+    for test_input in test_inputs:
+      with self.session() as sess:
+        with self.test_scope():
+          a = array_ops.placeholder(dtype)
+          index = array_ops.placeholder(dtypes.int32)
+          out = math_ops.reduce_sum(a, index)
+        result = sess.run(out, {
+            a: np.array(test_input, dtype=dtype),
+            index: [0]
+        })
+        # Compare the results using float32 type.
+        self.assertAllClose(
+            np.float32(result),
+            np.float32(expected_result),
+            rtol=rtol,
+            atol=atol)
+
+  def testReduceSumF16(self):
+    """Tests the reduce sum of float16 doesn't lose too much precision."""
+
+    if np.float16 not in self.all_types:
+      return
+
+    f16_max = np.finfo(np.float16).max
+    self._testReduceSum(
+        f16_max, np.float16,
+        itertools.permutations([f16_max, f16_max, f16_max * (-1.0)], 3))
+
+  def testReduceSumBF16(self):
+    """Tests the reduce sum of bfloat16 doesn't lose too much precision."""
+
+    if dtypes.bfloat16.as_numpy_dtype not in self.all_types:
+      return
+
+    bf16_max = np.float32(dtypes.bfloat16.max)
+    f32_max = dtypes.float32.max
+    value = min(bf16_max, f32_max - bf16_max) / 2
+    self._testReduceSum(
+        dtypes.bfloat16.as_numpy_dtype(value), dtypes.bfloat16.as_numpy_dtype,
+        itertools.permutations([bf16_max, value, bf16_max * (-1.0)], 3))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/reduce_window_test.py

@@ -0,0 +1,98 @@
+# Copyright 2018 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.reduce_window."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.compiler.tf2xla.python import xla
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import googletest
+
+
+class ReduceWindowTest(xla_test.XLATestCase):
+  """Test cases for xla.reduce_window."""
+
+  def _reduce_window(self, operand, init, reducer, **kwargs):
+    with self.session():
+      placeholder = array_ops.placeholder(operand.dtype)
+      with self.test_scope():
+        output = xla.reduce_window(placeholder, init, reducer, **kwargs)
+      return output.eval(feed_dict={placeholder: operand})
+
+  def testReduceWindow(self):
+
+    # TODO(b/77644762): float16 and float64 ReduceWindow are unimplemented.
+    for dtype in set(self.numeric_types).intersection(
+        set([dtypes.bfloat16.as_numpy_dtype, np.float32])):
+
+      @function.Defun(dtype, dtype)
+      def sum_reducer(x, y):
+        return x + y
+
+      @function.Defun(dtype, dtype)
+      def mul_reducer(x, y):
+        return x * y
+
+      self.assertAllClose(
+          np.array([3, 5, 7, 9, 11, 13], dtype=dtype),
+          self._reduce_window(
+              np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype),
+              0.0,
+              sum_reducer,
+              window_dimensions=[2]))
+
+      self.assertAllClose(
+          np.array([3, 7, 11], dtype=dtype),
+          self._reduce_window(
+              np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype),
+              0.0,
+              sum_reducer,
+              window_dimensions=[2],
+              window_strides=[2]))
+
+      self.assertAllClose(
+          np.array([1, 4, 7], dtype=dtype),
+          self._reduce_window(
+              np.array([1, 2, 3, 4, 5, 6, 7], dtype=dtype),
+              0.0,
+              sum_reducer,
+              window_dimensions=[1],
+              window_strides=[3]))
+
+      self.assertAllClose(
+          np.array([[24, 36, 24], [96, 0, 0]], dtype=dtype),
+          self._reduce_window(
+              np.array([[1, 2, 3, 4], [4, 3, 2, 1], [2, 4, 0, 1]], dtype=dtype),
+              1.0,
+              mul_reducer,
+              window_dimensions=[2, 2],
+              window_strides=[1, 1]))
+
+      self.assertAllClose(
+          np.array([[0, 0, 0], [5, 10, 5], [2, 4, 1], [0, 0, 0]], dtype=dtype),
+          self._reduce_window(
+              np.array([[1, 2, 3, 4], [4, 3, 2, 1], [2, 4, 0, 1]], dtype=dtype),
+              0.0,
+              sum_reducer,
+              window_dimensions=[2, 2],
+              window_strides=[2, 2],
+              padding=[[2, 3], [1, 2]]))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/repeat_op_test.py

@@ -0,0 +1,47 @@
+# Copyright 2021 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 tensorflow.ops.array_ops.repeat."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class RepeatTest(xla_test.XLATestCase):
+
+  def test(self):
+
+    # Verifies that bounded dynamic result generated from the Where op can be
+    # Reshaped correctly.
+    @def_function.function(jit_compile=True)
+    def repeat(values, repeats, axis):
+      return array_ops.repeat(values, repeats, axis)
+
+    with self.session() as sess:
+      with self.test_scope():
+        values = array_ops.constant([[1, 2], [3, 4]], dtype=dtypes.int32)
+        repeats = array_ops.constant([1, 2], dtype=dtypes.int32)
+        y1 = repeat(values, repeats, 0)
+        y2 = repeat(values, repeats, 1)
+      actual1, actual2 = sess.run([y1, y2])
+
+    self.assertAllEqual(actual1, [[1, 2], [3, 4], [3, 4]])
+    self.assertAllEqual(actual2, [[1, 2, 2], [3, 4, 4]])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/reshape_op_test.py

@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""Tests for reshape."""
+
+from absl.testing import parameterized
+
+from tensorflow.compiler.tests import xla_test
+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 ReshapeTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(('32_bit_index', dtypes.int32),
+                                  ('64_bit_index', dtypes.int64))
+  def testBasic(self, index_dtype):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[2, 3])
+        with self.test_scope():
+          shape = constant_op.constant([3, 2], dtype=index_dtype)
+          o = array_ops.reshape(i, shape)
+        params = {
+            i: [[1, 2, 3], [4, 5, 6]],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[1, 2], [3, 4], [5, 6]], result)
+
+  def testInt64(self):
+    with self.session():
+      with self.test_scope():
+        x = array_ops.zeros([50000, 50000], dtype=dtypes.bool)
+        # Provide dimension larger than int32
+        y = array_ops.reshape(x, [50000**2])
+        self.assertEqual([50000**2], y.get_shape().as_list())
+        # Even if first dimension is within int32, ensure we correctly go to
+        # int64
+        y = array_ops.reshape(x, [1, 50000**2])
+        self.assertEqual([1, 50000**2], y.get_shape().as_list())
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/compiler/tests/reverse_ops_test.py

@@ -0,0 +1,69 @@
+# 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."""
+
+import itertools
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+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(xla_test.XLATestCase):
+
+  def testReverseOneDim(self):
+    shape = (7, 5, 9, 11)
+    for revdim in range(-len(shape), len(shape)):
+      self._AssertReverseEqual([revdim], shape)
+
+  def testReverseMoreThanOneDim(self):
+    shape = (7, 5, 9, 11)
+    # The offset is used to test various (but not all) combinations of negative
+    # and positive axis indices that are guaranteed to not collide at the same
+    # index.
+    for revdims in itertools.chain.from_iterable(
+        itertools.combinations(range(-offset,
+                                     len(shape) - offset), k)
+        for k in range(2,
+                       len(shape) + 1)
+        for offset in range(0, len(shape))):
+      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.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 = tuple(
+            slice(-1, None, -1)
+            if d in revdims or d - len(shape) 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/tests/reverse_sequence_op_args_test.py

@@ -0,0 +1,52 @@
+# Copyright 2023 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 tensorflow.ops.reverse_sequence_op."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.compat import v2_compat
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import errors
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ReverseSequenceArgsTest(xla_test.XLATestCase):
+  """Tests argument verification of array_ops.reverse_sequence."""
+
+  def testInvalidArguments(self):
+    # seq_axis negative
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError), "seq_dim must be >=0"
+    ):
+
+      @def_function.function(jit_compile=True)
+      def f(x):
+        return array_ops.reverse_sequence(x, [2, 2], seq_axis=-1)
+
+      f([[1, 2], [3, 4]])
+
+    # batch_axis negative
+    with self.assertRaisesRegex(ValueError, "batch_dim must be >=0"):
+
+      @def_function.function(jit_compile=True)
+      def g(x):
+        return array_ops.reverse_sequence(x, [2, 2], seq_axis=1, batch_axis=-1)
+
+      g([[1, 2], [3, 4]])
+
+
+if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
+  test.main()

tensorflow/compiler/tests/reverse_sequence_op_test.py

@@ -0,0 +1,89 @@
+# Copyright 2015 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 tensorflow.ops.reverse_sequence_op."""
+
+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.platform import test
+
+
+class ReverseSequenceTest(xla_test.XLATestCase):
+
+  def _testReverseSequence(self,
+                           x,
+                           batch_axis,
+                           seq_axis,
+                           seq_lengths,
+                           truth,
+                           expected_err_re=None):
+    with self.session():
+      p = array_ops.placeholder(dtypes.as_dtype(x.dtype))
+      lengths = array_ops.placeholder(dtypes.as_dtype(seq_lengths.dtype))
+      with self.test_scope():
+        ans = array_ops.reverse_sequence(
+            p, batch_axis=batch_axis, seq_axis=seq_axis, seq_lengths=lengths)
+      if expected_err_re is None:
+        tf_ans = ans.eval(feed_dict={p: x, lengths: seq_lengths})
+        self.assertAllClose(tf_ans, truth, atol=1e-10)
+      else:
+        with self.assertRaisesOpError(expected_err_re):
+          ans.eval(feed_dict={p: x, lengths: seq_lengths})
+
+  def testSimple(self):
+    x = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.int32)
+    expected = np.array([[1, 2, 3], [6, 5, 4], [8, 7, 9]], dtype=np.int32)
+    self._testReverseSequence(
+        x,
+        batch_axis=0,
+        seq_axis=1,
+        seq_lengths=np.array([1, 3, 2], np.int32),
+        truth=expected)
+
+  def _testBasic(self, dtype, len_dtype):
+    x = np.asarray(
+        [[[1, 2, 3, 4], [5, 6, 7, 8]], [[9, 10, 11, 12], [13, 14, 15, 16]],
+         [[17, 18, 19, 20], [21, 22, 23, 24]]],
+        dtype=dtype)
+    x = x.reshape(3, 2, 4, 1, 1)
+    x = x.transpose([2, 1, 0, 3, 4])  # permute axes 0 <=> 2
+
+    # reverse dim 2 up to (0:3, none, 0:4) along dim=0
+    seq_lengths = np.asarray([3, 0, 4], dtype=len_dtype)
+
+    truth_orig = np.asarray(
+        [
+            [[3, 2, 1, 4], [7, 6, 5, 8]],  # reverse 0:3
+            [[9, 10, 11, 12], [13, 14, 15, 16]],  # reverse none
+            [[20, 19, 18, 17], [24, 23, 22, 21]]
+        ],  # reverse 0:4 (all)
+        dtype=dtype)
+    truth_orig = truth_orig.reshape(3, 2, 4, 1, 1)
+    truth = truth_orig.transpose([2, 1, 0, 3, 4])  # permute axes 0 <=> 2
+
+    seq_axis = 0  # permute seq_axis and batch_axis (originally 2 and 0, resp.)
+    batch_axis = 2
+    self._testReverseSequence(x, batch_axis, seq_axis, seq_lengths, truth)
+
+  def testSeqLength(self):
+    for dtype in self.all_types:
+      for seq_dtype in self.all_types & {np.int32, np.int64}:
+        self._testBasic(dtype, seq_dtype)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/rmsprop_test.py

@@ -0,0 +1,128 @@
+# 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.
+# ==============================================================================
+"""Tests for RMSProp optimizer."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+from tensorflow.python.training import rmsprop
+
+
+class RmspropTest(xla_test.XLATestCase):
+
+  def _rmsprop_update_numpy(self,
+                            var,
+                            g,
+                            mg,
+                            rms,
+                            mom,
+                            lr,
+                            decay=0.9,
+                            momentum=0.0,
+                            epsilon=1e-10,
+                            centered=False):
+    rms_t = rms * decay + (1 - decay) * g * g
+    denom_t = rms_t + epsilon
+    if centered:
+      mg_t = mg * decay + (1 - decay) * g
+      denom_t -= mg_t * mg_t
+    else:
+      mg_t = mg
+    mom_t = momentum * mom + lr * g / np.sqrt(denom_t, dtype=denom_t.dtype)
+    var_t = var - mom_t
+    return var_t, mg_t, rms_t, mom_t
+
+  def testBasic(self):
+    for dtype in self.float_types | self.complex_types:
+      for centered in [False, True]:
+        with self.session(), self.test_scope():
+          # Initialize variables for numpy implementation.
+          var0_np = np.array([1.0, 2.0], dtype=dtype)
+          grads0_np = np.array([0.1, 0.1], dtype=dtype)
+          var1_np = np.array([3.0, 4.0], dtype=dtype)
+          grads1_np = np.array([0.01, 0.01], dtype=dtype)
+          mg0_np = np.array([0.0, 0.0], dtype=dtype)
+          mg1_np = np.array([0.0, 0.0], dtype=dtype)
+          rms0_np = np.array([1.0, 1.0], dtype=dtype)
+          rms1_np = np.array([1.0, 1.0], dtype=dtype)
+          mom0_np = np.array([0.0, 0.0], dtype=dtype)
+          mom1_np = np.array([0.0, 0.0], dtype=dtype)
+
+          var0 = resource_variable_ops.ResourceVariable(var0_np)
+          var1 = resource_variable_ops.ResourceVariable(var1_np)
+          grads0 = constant_op.constant(grads0_np)
+          grads1 = constant_op.constant(grads1_np)
+          learning_rate = 3.0
+          rms_opt = rmsprop.RMSPropOptimizer(learning_rate, centered=centered)
+          rms_update = rms_opt.apply_gradients(
+              zip([grads0, grads1], [var0, var1]))
+          self.evaluate(variables.global_variables_initializer())
+
+          mg0 = rms_opt.get_slot(var0, "mg")
+          self.assertEqual(mg0 is not None, centered)
+          mg1 = rms_opt.get_slot(var1, "mg")
+          self.assertEqual(mg1 is not None, centered)
+          rms0 = rms_opt.get_slot(var0, "rms")
+          self.assertIsNotNone(rms0)
+          rms1 = rms_opt.get_slot(var1, "rms")
+          self.assertIsNotNone(rms1)
+          mom0 = rms_opt.get_slot(var0, "momentum")
+          self.assertIsNotNone(mom0)
+          mom1 = rms_opt.get_slot(var1, "momentum")
+          self.assertIsNotNone(mom1)
+
+          # Fetch params to validate initial values
+          self.assertAllClose([1.0, 2.0], self.evaluate(var0))
+          self.assertAllClose([3.0, 4.0], self.evaluate(var1))
+
+          # Run 3 steps of RMSProp
+          for _ in range(3):
+            self.evaluate(rms_update)
+
+            var0_np, mg0_np, rms0_np, mom0_np = self._rmsprop_update_numpy(
+                var0_np,
+                grads0_np,
+                mg0_np,
+                rms0_np,
+                mom0_np,
+                learning_rate,
+                centered=centered)
+            var1_np, mg1_np, rms1_np, mom1_np = self._rmsprop_update_numpy(
+                var1_np,
+                grads1_np,
+                mg1_np,
+                rms1_np,
+                mom1_np,
+                learning_rate,
+                centered=centered)
+
+            # Validate updated params
+            if centered:
+              self.assertAllCloseAccordingToType(mg0_np, self.evaluate(mg0))
+              self.assertAllCloseAccordingToType(mg1_np, self.evaluate(mg1))
+            self.assertAllCloseAccordingToType(rms0_np, self.evaluate(rms0))
+            self.assertAllCloseAccordingToType(rms1_np, self.evaluate(rms1))
+            self.assertAllCloseAccordingToType(mom0_np, self.evaluate(mom0))
+            self.assertAllCloseAccordingToType(mom1_np, self.evaluate(mom1))
+            self.assertAllCloseAccordingToType(var0_np, self.evaluate(var0))
+            self.assertAllCloseAccordingToType(var1_np, self.evaluate(var1))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/runtime_shape_check_test.py

@@ -0,0 +1,70 @@
+# Copyright 2022 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 shape checks at runtime in XLA:GPU."""
+
+from tensorflow.compiler.tests import xla_test
+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 errors
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class RuntimeShapeCheckTest(xla_test.XLATestCase):
+
+  def testUniqueDifferentSizes(self):
+    """Test that we correctly check for shape mismatches at runtime."""
+    if 'tpu' in self.device.lower():
+      self.skipTest('We do not check shapes on TPU')
+
+    with ops.device(f'device:{self.device}:0'):
+
+      @def_function.function(jit_compile=True)
+      def f(x, y):
+        return array_ops.unique(x).y + array_ops.unique(y).y
+
+      f(constant_op.constant([3.1, 3.2]), constant_op.constant([3.3, 3.2]))
+
+      with self.assertRaisesRegex(errors.InternalError, 'different size'):
+        f(
+            constant_op.constant([3.1, 3.2]),
+            constant_op.constant([3.1, 3.2, 3.3]))
+
+  def testWhereOpDifferentSizes(self):
+    """Test shape mismatches with multiple dimensions."""
+    if 'tpu' in self.device.lower():
+      self.skipTest('We do not check shapes on TPU')
+
+    with ops.device(f'device:{self.device}:0'):
+
+      @def_function.function(jit_compile=True)
+      def f(x, y):
+        return array_ops.where(x) + array_ops.where(y)
+
+      f(
+          constant_op.constant([[3.1, 3.2, 0], [3.1, 3.2, 0]]),
+          constant_op.constant([[3.3, 3.2, 0, 0, 0], [3.3, 3.2, 0, 0, 0]]))
+
+      with self.assertRaisesRegex(errors.InternalError, 'different size'):
+        f(
+            constant_op.constant([[3.1, 3.2, 0], [3.1, 3.2, 0]]),
+            constant_op.constant([[3.3, 3.2, 0, 0, 0], [3.3, 3.2, 3.3, 0, 0]]))
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

tensorflow/compiler/tests/scan_ops_test.py

@@ -0,0 +1,319 @@
+# Copyright 2016 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 scan ops."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+def numpy_reverse(x, axis):
+  length = len(x.shape)
+  if axis < 0:
+    axis = length + axis
+
+  ix = tuple(
+      slice(None, None, -1) if i == axis else slice(None) for i in range(length)
+  )
+  return x[ix]
+
+
+def handle_options(func, init_fn, x, axis, exclusive, reverse):
+  """Adds tf options to numpy scan ops."""
+  length = len(x.shape)
+  if axis < 0:
+    axis = length + axis
+
+  if reverse:
+    x = numpy_reverse(x, axis)
+
+  if exclusive:
+    ix_head = tuple(slice(0, 1) if i == axis else slice(None)
+                    for i in range(length))
+    ix_init = tuple(
+        slice(0, -1) if i == axis else slice(None) for i in range(length)
+    )
+    init = init_fn(x[ix_head])
+    x = np.concatenate([init, func(x[ix_init], axis=axis)], axis=axis)
+  else:
+    x = func(x, axis=axis)
+
+  if reverse:
+    x = numpy_reverse(x, axis)
+  return x
+
+
+class CumsumTest(xla_test.XLATestCase):
+
+  valid_dtypes = [np.float32, np.int32, np.int64]
+
+  def axis_dtypes(self):
+    return set(self.int_types).intersection([np.int32, np.int64])
+
+  def _compare(self, x, axis, exclusive, reverse):
+    np_out = handle_options(np.cumsum, np.zeros_like, x, axis, exclusive,
+                            reverse)
+    with self.session(), self.test_scope():
+      p = array_ops.placeholder(x.dtype)
+      tf_out = math_ops.cumsum(p, axis, exclusive, reverse).eval(
+          feed_dict={p: x})
+
+    self.assertAllClose(np_out, tf_out)
+
+  def _compareAll(self, x, axis):
+    for exclusive in [True, False]:
+      for reverse in [True, False]:
+        self._compare(x, axis, exclusive, reverse)
+
+  def testEmpty(self):
+    for dtype in self.valid_dtypes:
+      x = np.zeros([0]).astype(dtype)
+      for axis in (-1, 0):
+        self._compareAll(x, axis)
+
+  def testAxisType(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 6).reshape([5]).astype(dtype)
+      for axis_dtype in self.axis_dtypes():
+        with self.session(), self.test_scope():
+          p = array_ops.placeholder(x.dtype)
+          axis = constant_op.constant(0, axis_dtype)
+          math_ops.cumsum(p, axis).eval(feed_dict={p: x})
+
+  def test1D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 6).reshape([5]).astype(dtype)
+      for axis in (-1, 0):
+        self._compareAll(x, axis)
+
+  def test2D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(0, 10).reshape([2, 5]).astype(dtype)
+      for axis in (-2, -1, 0, 1):
+        self._compareAll(x, axis)
+
+  def test3D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(0, 20).reshape([2, 2, 5]).astype(dtype)
+      for axis in (-3, -2, -1, 0, 1, 2):
+        self._compareAll(x, axis)
+
+  def test6D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 145).reshape([2, 2, 3, 3, 2, 2]).astype(dtype)
+      for axis in range(-6, 6, 3):
+        self._compareAll(x, axis)
+
+  def testMixedPrecision(self):
+    with self.session(), self.test_scope():
+      y = math_ops.cumsum(
+          constant_op.constant([1., 2., 3., 4.], dtypes.bfloat16),
+          -1,
+          exclusive=True).eval()
+    self.assertAllEqual(y, [0., 1., 3., 6.])
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testInvalidAxis(self):
+    x = np.arange(0, 10).reshape([2, 5]).astype(np.float32)
+    with self.session(), self.test_scope():
+      input_tensor = ops.convert_to_tensor(x)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "Expected scan axis in the range [-2, 2)" in str(e)):
+        math_ops.cumsum(input_tensor, -3).eval()
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "Expected scan axis in the range [-2, 2)" in str(e)):
+        math_ops.cumsum(input_tensor, 2).eval()
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "axis must be a scalar" in str(e)):
+        math_ops.cumsum(input_tensor, [0]).eval()
+
+
+class CumulativeLogSumExpTest(xla_test.XLATestCase):
+
+  valid_dtypes = [np.float32, np.float64]
+
+  def axis_dtypes(self):
+    return set(self.int_types).intersection([np.int32, np.int64])
+
+  def _compare(self, x, axis, exclusive, reverse):
+
+    def neginf_like(x):
+      return -np.inf * np.ones_like(x)
+
+    np_out = handle_options(np.logaddexp.accumulate, neginf_like, x, axis,
+                            exclusive, reverse)
+    with self.session(), self.test_scope():
+      p = array_ops.placeholder(x.dtype)
+      tf_out = math_ops.cumulative_logsumexp(p, axis, exclusive,
+                                             reverse).eval(feed_dict={p: x})
+
+    self.assertAllClose(np_out, tf_out, rtol=4e-5)
+
+  def _compareAll(self, x, axis):
+    for exclusive in [True, False]:
+      for reverse in [True, False]:
+        self._compare(x, axis, exclusive, reverse)
+
+  def testEmpty(self):
+    for dtype in self.valid_dtypes:
+      x = np.zeros([0]).astype(dtype)
+      for axis in (-1, 0):
+        self._compareAll(x, axis)
+
+  def testAxisType(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 6).reshape([5]).astype(dtype)
+      for axis_dtype in self.axis_dtypes():
+        with self.session(), self.test_scope():
+          p = array_ops.placeholder(x.dtype)
+          axis = constant_op.constant(0, axis_dtype)
+          math_ops.cumulative_logsumexp(p, axis).eval(feed_dict={p: x})
+
+  def test1D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 6).reshape([5]).astype(dtype)
+      for axis in (-1, 0):
+        self._compareAll(x, axis)
+
+  def test2D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(0, 10).reshape([2, 5]).astype(dtype)
+      for axis in (-2, -1, 0, 1):
+        self._compareAll(x, axis)
+
+  def test3D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(0, 20).reshape([2, 2, 5]).astype(dtype)
+      for axis in (-3, -2, -1, 0, 1, 2):
+        self._compareAll(x, axis)
+
+  def test6D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 145).reshape([2, 2, 3, 3, 2, 2]).astype(dtype)
+      for axis in range(-6, 6, 3):
+        self._compareAll(x, axis)
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testInvalidAxis(self):
+    x = np.arange(0, 10).reshape([2, 5]).astype(np.float32)
+    with self.session(), self.test_scope():
+      input_tensor = ops.convert_to_tensor(x)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "Expected scan axis in the range [-2, 2)" in str(e)):
+        math_ops.cumulative_logsumexp(input_tensor, -3).eval()
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "Expected scan axis in the range [-2, 2)" in str(e)):
+        math_ops.cumulative_logsumexp(input_tensor, 2).eval()
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "axis must be a scalar" in str(e)):
+        math_ops.cumulative_logsumexp(input_tensor, [0]).eval()
+
+
+class CumprodTest(xla_test.XLATestCase):
+
+  valid_dtypes = [np.float32, np.int32]
+
+  def axis_dtypes(self):
+    return set(self.int_types).intersection([np.int32, np.int64])
+
+  def _compare(self, x, axis, exclusive, reverse):
+    np_out = handle_options(np.cumprod, np.ones_like, x, axis, exclusive,
+                            reverse)
+    with self.session(), self.test_scope():
+      p = array_ops.placeholder(x.dtype)
+      prod = math_ops.cumprod(p, axis, exclusive, reverse)
+      tf_out = prod.eval(feed_dict={p: x})
+
+    self.assertAllClose(np_out, tf_out)
+
+  def _compareAll(self, x, axis):
+    for exclusive in [True, False]:
+      for reverse in [True, False]:
+        self._compare(x, axis, exclusive, reverse)
+
+  def testEmpty(self):
+    for dtype in self.valid_dtypes:
+      x = np.zeros([0]).astype(dtype)
+      for axis in (-1, 0):
+        self._compareAll(x, axis)
+
+  def testAxisType(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 6).reshape([5]).astype(dtype)
+      for axis_dtype in self.axis_dtypes():
+        with self.session(), self.test_scope():
+          p = array_ops.placeholder(x.dtype)
+          axis = constant_op.constant(0, axis_dtype)
+          math_ops.cumprod(x, axis).eval(feed_dict={p: x})
+
+  def test1D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 6).reshape([5]).astype(dtype)
+      for axis in (-1, 0):
+        self._compareAll(x, axis)
+
+  def test2D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 11).reshape([2, 5]).astype(dtype)
+      for axis in (-2, -1, 0, 1):
+        self._compareAll(x, axis)
+
+  def test3D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 21).reshape([2, 2, 5]).astype(dtype)
+      for axis in (-3, -2, -1, 0, 1, 2):
+        self._compareAll(x, axis)
+
+  def test6D(self):
+    for dtype in self.valid_dtypes:
+      x = np.arange(1, 145).reshape([2, 2, 3, 3, 2, 2]).astype(dtype)
+      for axis in range(-6, 6, 3):
+        self._compareAll(x, axis)
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testInvalidAxis(self):
+    x = np.arange(0, 10).reshape([2, 5]).astype(np.float32)
+    with self.session(), self.test_scope():
+      input_tensor = ops.convert_to_tensor(x)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "Expected scan axis in the range [-2, 2)" in str(e)):
+        math_ops.cumprod(input_tensor, -3).eval()
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "Expected scan axis in the range [-2, 2)" in str(e)):
+        math_ops.cumprod(input_tensor, 2).eval()
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          lambda e: "axis must be a scalar" in str(e)):
+        math_ops.cumprod(input_tensor, [0]).eval()
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/scatter_nd_op_test.py

@@ -0,0 +1,245 @@
+# Copyright 2018 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 tensorflow.ops.tf.scatter_nd."""
+
+import functools
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+def _AsType(v, vtype):
+  return v.astype(vtype) if isinstance(v, np.ndarray) else vtype(v)
+
+
+def _FlatInnerDims(tensor, ndims=2):
+  shape = list(tensor.shape)
+  return tensor.reshape(
+      [functools.reduce(lambda x, y: x * y, shape[:-ndims + 1], 1)] +
+      shape[-ndims + 1:])
+
+
+def _FlatOuterDims(tensor, ndims=2):
+  shape = list(tensor.shape)
+  return tensor.reshape(
+      shape[:ndims - 1] +
+      [functools.reduce(lambda x, y: x * y, shape[ndims - 1:], 1)])
+
+
+def _NumpyScatterNd(ref, indices, updates, op):
+  ixdim = indices.shape[-1]
+  num_updates = indices.size // ixdim
+  total_nd = len(ref.shape)
+  slice_size = 1
+  for i in range(ixdim, total_nd):
+    slice_size *= ref.shape[i]
+  flat_indices = _FlatInnerDims(indices)
+  flat_updates = updates.reshape((num_updates, slice_size))
+  output_flat = _FlatOuterDims(ref, ixdim + 1)
+  for ix_updates, ix_output in enumerate(flat_indices):
+    ix_output = tuple(ix_output)
+    output_flat[ix_output] = op(output_flat[ix_output],
+                                flat_updates[ix_updates])
+  return output_flat.reshape(ref.shape)
+
+
+def _NumpyUpdate(indices, updates, shape):
+  ref = np.zeros(shape, dtype=updates.dtype)
+  return _NumpyScatterNd(ref, indices, updates, lambda p, u: u)
+
+
+class ScatterNdTest(xla_test.XLATestCase):
+
+  def _VariableRankTest(self,
+                        np_scatter,
+                        tf_scatter,
+                        vtype,
+                        itype,
+                        repeat_indices=False):
+    np.random.seed(8)
+    ref_shapes = [(3, 6), (3, 6), (3, 6, 9), (3, 6, 9), (3, 6, 9), (3, 6, 9)]
+    indices_shapes = [(2,), (2, 2), (2,), (2, 2), (2, 3), (2, 3, 3)]
+    for ref_shape, indices_shape in zip(ref_shapes, indices_shapes):
+      num_updates = indices_shape[0]
+      ixdim = indices_shape[-1]
+
+      indexable_area_shape = ()
+      for i in range(ixdim):
+        indexable_area_shape += (ref_shape[i],)
+      all_indices = [
+          list(coord)
+          for coord, _ in np.ndenumerate(np.empty(indexable_area_shape, vtype))
+      ]
+      np.random.shuffle(all_indices)
+      indices = np.array(all_indices[:num_updates])
+
+      if num_updates > 1 and repeat_indices:
+        indices = indices[:num_updates // 2]
+        for _ in range(num_updates - num_updates // 2):
+          indices = np.append(
+              indices, [indices[np.random.randint(num_updates // 2)]], axis=0)
+        np.random.shuffle(indices)
+      indices = _AsType(indices[:num_updates], itype)
+
+      updates_shape = (num_updates,)
+      for i in range(ixdim, len(ref_shape)):
+        updates_shape += (ref_shape[i],)
+      updates = _AsType(np.random.randn(*(updates_shape)), vtype)
+
+      # Scatter via numpy
+      np_out = np_scatter(indices, updates, ref_shape)
+      # Scatter via tensorflow
+      tf_out = tf_scatter(indices, updates, ref_shape)
+
+      self.assertAllClose(np_out, tf_out)
+
+  def _VariableRankTests(self, np_scatter, tf_scatter):
+    for vtype in self.numeric_types:
+      for itype in set([np.int32, np.int64]).intersection(set(self.int_types)):
+        self._VariableRankTest(np_scatter, tf_scatter, vtype, itype)
+
+  def _runScatterNd(self, indices, updates, shape):
+    with self.session():
+      updates_placeholder = array_ops.placeholder(updates.dtype)
+      indices_placeholder = array_ops.placeholder(indices.dtype)
+      with self.test_scope():
+        output = array_ops.scatter_nd(indices_placeholder, updates_placeholder,
+                                      shape)
+      feed_dict = {updates_placeholder: updates, indices_placeholder: indices}
+      return output.eval(feed_dict=feed_dict)
+
+  def testSimple(self):
+    indices = np.array([[4], [3], [1], [7]], dtype=np.int32)
+    updates = np.array([9, 10, 11, 12], dtype=np.float32)
+    expected = np.array([0, 11, 0, 10, 9, 0, 0, 12], dtype=np.int32)
+    self.assertAllEqual(expected, self._runScatterNd(indices, updates, [8]))
+
+  def testRepeatedIndices(self):
+    indices = np.array([[0], [1], [0], [1]], dtype=np.int32)
+    updates = np.array([9, 10, 11, 12], dtype=np.float32)
+    expected = np.array([20, 22], dtype=np.int32)
+    self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2]))
+
+  def testSimple2(self):
+    indices = np.array([[1, 0], [1, 1]], dtype=np.int32)
+    updates = np.array([11., 12.], dtype=np.float32)
+    expected = np.array([[0., 0.], [11., 12.], [0., 0.]], dtype=np.float32)
+    self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2]))
+
+  def testSimple3(self):
+    indices = np.array([[1]], dtype=np.int32)
+    updates = np.array([[11., 12.]], dtype=np.float32)
+    expected = np.array([[0., 0.], [11., 12.], [0., 0.]])
+    self.assertAllEqual(expected, self._runScatterNd(indices, updates, [3, 2]))
+
+  def testVariableRankUpdate(self):
+    self._VariableRankTests(_NumpyUpdate, self._runScatterNd)
+
+  def testExtraIndicesDimensions(self):
+    indices = np.zeros([1, 1, 2], np.int32)
+    updates = np.zeros([1, 1], np.int32)
+    expected = np.zeros([2, 2], dtype=np.int32)
+    self.assertAllEqual(expected, self._runScatterNd(indices, updates, [2, 2]))
+
+  @test_util.disable_mlir_bridge("Error messages differ")
+  def testRank3InvalidShape1(self):
+    indices = np.zeros([3, 2, 2], np.int32)
+    updates = np.zeros([2, 2, 2], np.int32)
+    with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
+                                             "Must have updates.shape"):
+      self._runScatterNd(indices, updates, [2, 2, 2])
+
+  @test_util.disable_mlir_bridge("Error messages differ")
+  def testRank3InvalidShape2(self):
+    indices = np.zeros([2, 2, 1], np.int32)
+    updates = np.zeros([2, 2], np.int32)
+    with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
+                                             "Must have updates.shape"):
+      self._runScatterNd(indices, updates, [2, 2, 2])
+
+  def testScatterOutOfRange(self):
+    updates = np.array([-3, -4, -5]).astype(np.float32)
+
+    # Indices all in range, no problem.
+    indices = np.array([[2], [0], [5]], dtype=np.int32)
+    self._runScatterNd(indices, updates, [6])
+
+    # Indices out of range should not fail. It produces implementation-defined
+    # output.
+    indices = np.array([[-1], [0], [5]], dtype=np.int32)
+    self._runScatterNd(indices, updates, [6])
+    indices = np.array([[2], [0], [6]], dtype=np.int32)
+    self._runScatterNd(indices, updates, [6])
+
+
+class ScatterNdTensorTest(xla_test.XLATestCase):
+
+  def _runScatter(self, op):
+    indices_np = np.array([[4], [3], [1], [7]], dtype=np.int32)
+    updates_np = np.array([9, 10, 11, 12], dtype=np.float32)
+    with self.session() as sess, self.test_scope():
+      indices = array_ops.placeholder(indices_np.dtype, shape=indices_np.shape)
+      updates = array_ops.placeholder(updates_np.dtype, shape=updates_np.shape)
+      t = array_ops.ones([8], dtype=np.float32)
+
+      out = op(t, indices, updates)
+      return sess.run(out, feed_dict={indices: indices_np, updates: updates_np})
+
+  def testAdd(self):
+    self.assertAllEqual(
+        self._runScatter(array_ops.tensor_scatter_add),
+        np.array([1, 12, 1, 11, 10, 1, 1, 13], dtype=np.float32))
+
+  def testSub(self):
+    self.assertAllEqual(
+        self._runScatter(array_ops.tensor_scatter_sub),
+        np.array([1, -10, 1, -9, -8, 1, 1, -11], dtype=np.float32))
+
+  def testUpdate(self):
+    self.assertAllEqual(
+        self._runScatter(array_ops.tensor_scatter_update),
+        np.array([1, 11, 1, 10, 9, 1, 1, 12], dtype=np.float32))
+
+
+class ScatterNdTensorScalarUpdateTest(xla_test.XLATestCase):
+
+  def _runScatter(self, op):
+    indices_np = np.array([[4], [3], [1], [7]], dtype=np.int32)
+    updates_np = np.array(9, dtype=np.float32)
+    with self.session() as sess, self.test_scope():
+      indices = array_ops.placeholder(indices_np.dtype, shape=indices_np.shape)
+      updates = array_ops.placeholder(updates_np.dtype, shape=updates_np.shape)
+      t = array_ops.ones([8], dtype=np.float32)
+
+      out = op(t, indices, updates)
+      return sess.run(out, feed_dict={indices: indices_np, updates: updates_np})
+
+  def testUpdate(self):
+    self.assertAllEqual(
+        self._runScatter(array_ops.tensor_scatter_update),
+        np.array([1, 9, 1, 9, 9, 1, 1, 9], dtype=np.float32))
+
+  def testAdd(self):
+    self.assertAllEqual(
+        self._runScatter(array_ops.tensor_scatter_add),
+        np.array([1, 10, 1, 10, 10, 1, 1, 10], dtype=np.float32))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/searchsorted_op_test.py

@@ -0,0 +1,92 @@
+# 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.
+# ==============================================================================
+"""Test for XLA implementation of tf.searchsorted."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class SearchSorteddOpTest(xla_test.XLATestCase):
+
+  def test1D(self):
+    # Test against NumPy implementation (which is 1D only).
+    np.random.seed(1)
+    for side in ['left', 'right']:
+      for dtype in [np.float32, np.int32]:
+        values = np.random.uniform(
+            low=-1000, high=1000, size=(10,)).astype(dtype)
+        unsorted = np.random.uniform(
+            low=-1000, high=1000, size=(20,)).astype(dtype)
+
+        sorted_sequence = np.sort(unsorted)
+        np_ans = np.searchsorted(sorted_sequence, values, side=side)
+
+        with self.session() as session:
+          with self.test_scope():
+            tf_ans = array_ops.searchsorted(sorted_sequence, values, side=side)
+          tf_out = session.run(tf_ans)
+          self.assertAllEqual(np_ans, tf_out)
+
+  def testNanValue(self):
+    for dtype in self.float_types:
+      sorted_sequence = np.array([[1, 3, 5, 7, 9]], dtype)
+      values = np.array([[np.nan]], dtype)
+
+      self._test2DExample(
+          dtype,
+          'left',
+          sorted_sequence,
+          values,
+          np.array([[0]], dtype=np.int32),
+      )
+
+      self._test2DExample(
+          dtype,
+          'right',
+          sorted_sequence,
+          values,
+          np.array([[5]], dtype=np.int32),
+      )
+
+  def _test2DExample(self, dtype, side, sorted_sequence, values, correct_ans):
+
+    with self.session() as session:
+      with self.test_scope():
+        tf_ans = array_ops.searchsorted(sorted_sequence, values, side=side)
+      tf_out = session.run(tf_ans)
+      self.assertAllEqual(correct_ans, tf_out)
+
+  def testLowerBound2DExample(self):
+    # 2D TensorFlow documentation example.
+    for dtype in self.float_types | self.int_types:
+      sorted_sequence = np.array([[0, 3, 9, 9, 10], [1, 2, 3, 4, 5]], dtype)
+      values = np.array([[2, 4, 9], [0, 2, 6]], dtype)
+      correct_ans = np.array([[1, 2, 2], [0, 1, 5]], dtype)
+      self._test2DExample(dtype, 'left', sorted_sequence, values, correct_ans)
+
+  def testUpperBound2DExample(self):
+    # 2D TensorFlow documentation example.
+    for dtype in self.float_types | self.int_types:
+      sorted_sequence = np.array([[0, 3, 9, 9, 10], [1, 2, 3, 4, 5]], dtype)
+      values = np.array([[2, 4, 9], [0, 2, 6]], dtype)
+      correct_ans = np.array([[1, 2, 4], [0, 2, 5]], dtype)
+      self._test2DExample(dtype, 'right', sorted_sequence, values, correct_ans)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/tests/segment_reduction_ops_test.py

@@ -0,0 +1,338 @@
+# 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.
+# ==============================================================================
+"""Test cases for segment reduction ops."""
+
+import functools
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.client import device_lib
+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 SegmentReductionOpsTest(xla_test.XLATestCase):
+  """Test cases for segment reduction ops."""
+
+  def _findDevice(self, device_name):
+    devices = device_lib.list_local_devices()
+    for d in devices:
+      if d.device_type == device_name:
+        return True
+    return False
+
+  def _segmentReduction(self, op, data, indices, num_segments):
+    with self.session() as sess, self.test_scope():
+      d = array_ops.placeholder(data.dtype, shape=data.shape)
+      if isinstance(indices, int):
+        i = array_ops.placeholder(np.int32, shape=[])
+      else:
+        i = array_ops.placeholder(indices.dtype, shape=indices.shape)
+      return sess.run(op(d, i, num_segments), {d: data, i: indices})
+
+  def _unsortedSegmentSum(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.unsorted_segment_sum, data, indices,
+                                  num_segments)
+
+  def _segmentSumV2(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.segment_sum_v2, data, indices,
+                                  num_segments)
+
+  def _segmentProdV2(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.segment_prod_v2, data, indices,
+                                  num_segments)
+
+  def _segmentMinV2(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.segment_min_v2, data, indices,
+                                  num_segments)
+
+  def _segmentMaxV2(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.segment_max_v2, data, indices,
+                                  num_segments)
+
+  def _unsortedSegmentProd(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.unsorted_segment_prod, data, indices,
+                                  num_segments)
+
+  def _unsortedSegmentMin(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.unsorted_segment_min, data, indices,
+                                  num_segments)
+
+  def _unsortedSegmentMax(self, data, indices, num_segments):
+    return self._segmentReduction(math_ops.unsorted_segment_max, data, indices,
+                                  num_segments)
+
+  def testSegmentSum(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([1, 0, 2, 12], dtype=dtype),
+          self._segmentSumV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 4))
+
+  def testSegmentProd(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([0, 1, 2, 60], dtype=dtype),
+          self._segmentProdV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 4))
+
+  def testSegmentProdNumSegmentsLess(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([0, 1, 2], dtype=dtype),
+          self._segmentProdV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 3))
+
+  def testSegmentProdNumSegmentsMore(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([0, 1, 2, 60, 1], dtype=dtype),
+          self._segmentProdV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 5))
+
+  def testSegmentMin(self):
+    for dtype in self.int_types | self.float_types:
+      maxval = dtypes.as_dtype(dtype).max
+      if dtype == np.float64 and self._findDevice("TPU"):
+        maxval = np.inf
+      self.assertAllClose(
+          np.array([0, maxval, 2, 3], dtype=dtype),
+          self._segmentMinV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 4))
+
+  def testSegmentMinNumSegmentsLess(self):
+    for dtype in self.int_types | self.float_types:
+      maxval = dtypes.as_dtype(dtype).max
+      if dtype == np.float64 and self._findDevice("TPU"):
+        maxval = np.inf
+      self.assertAllClose(
+          np.array([0, maxval, 2], dtype=dtype),
+          self._segmentMinV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 3))
+
+  def testSegmentMinNumSegmentsMore(self):
+    for dtype in self.int_types | self.float_types:
+      maxval = dtypes.as_dtype(dtype).max
+      if dtype == np.float64 and self._findDevice("TPU"):
+        maxval = np.inf
+      self.assertAllClose(
+          np.array([0, maxval, 2, 3, maxval], dtype=dtype),
+          self._segmentMinV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 5))
+
+  def testSegmentMax(self):
+    for dtype in self.int_types | self.float_types:
+      minval = dtypes.as_dtype(dtype).min
+      if dtype == np.float64 and self._findDevice("TPU"):
+        minval = -np.inf
+      self.assertAllClose(
+          np.array([1, minval, 2, 5], dtype=dtype),
+          self._segmentMaxV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 4))
+
+  def testSegmentMaxNumSegmentsLess(self):
+    for dtype in self.int_types | self.float_types:
+      minval = dtypes.as_dtype(dtype).min
+      if dtype == np.float64 and self._findDevice("TPU"):
+        minval = -np.inf
+      self.assertAllClose(
+          np.array([1, minval, 2], dtype=dtype),
+          self._segmentMaxV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 3))
+
+  def testSegmentMaxNumSegmentsMore(self):
+    for dtype in self.int_types | self.float_types:
+      minval = dtypes.as_dtype(dtype).min
+      if dtype == np.float64 and self._findDevice("TPU"):
+        minval = -np.inf
+      self.assertAllClose(
+          np.array([1, minval, 2, 5, minval], dtype=dtype),
+          self._segmentMaxV2(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([0, 0, 2, 3, 3, 3], dtype=np.int32), 5))
+
+  def testUnsortedSegmentSum0DIndices1DData(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array(
+              [[0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0], [0, 1, 2, 3, 4, 5],
+               [0, 0, 0, 0, 0, 0]],
+              dtype=dtype),
+          self._unsortedSegmentSum(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype), 2, 4))
+
+  def testUnsortedSegmentSum1DIndices1DData(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([1, 3, 2, 9], dtype=dtype),
+          self._unsortedSegmentSum(
+              np.array([0, 1, 2, 3, 4, 5], dtype=dtype),
+              np.array([3, 0, 2, 1, 3, 3], dtype=np.int32), 4))
+
+  def testUnsortedSegmentSum1DIndices1DDataNegativeIndices(self):
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([6, 3, 0, 6], dtype=dtype),
+          self._unsortedSegmentSum(
+              np.array([0, 1, 2, 3, 4, 5, 6], dtype=dtype),
+              np.array([3, -1, 0, 1, 0, -1, 3], dtype=np.int32), 4))
+
+  def testUnsortedSegmentSum1DIndices2DDataDisjoint(self):
+    for dtype in self.numeric_types:
+      data = np.array(
+          [[0, 1, 2, 3], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43],
+           [50, 51, 52, 53]],
+          dtype=dtype)
+      indices = np.array([8, 1, 0, 3, 7], dtype=np.int32)
+      num_segments = 10
+      y = self._unsortedSegmentSum(data, indices, num_segments)
+      self.assertAllClose(
+          np.array(
+              [[30, 31, 32, 33], [20, 21, 22, 23], [0, 0, 0, 0],
+               [40, 41, 42, 43], [0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0],
+               [50, 51, 52, 53], [0, 1, 2, 3], [0, 0, 0, 0]],
+              dtype=dtype), y)
+
+  def testUnsortedSegmentSum1DIndices2DDataNonDisjoint(self):
+    for dtype in self.numeric_types:
+      data = np.array(
+          [[0, 1, 2, 3], [20, 21, 22, 23], [30, 31, 32, 33], [40, 41, 42, 43],
+           [50, 51, 52, 53]],
+          dtype=dtype)
+      indices = np.array([0, 1, 2, 0, 1], dtype=np.int32)
+      num_segments = 4
+      y = self._unsortedSegmentSum(data, indices, num_segments)
+      self.assertAllClose(
+          np.array(
+              [[40, 42, 44, 46], [70, 72, 74, 76], [30, 31, 32, 33],
+               [0, 0, 0, 0]],
+              dtype=dtype), y)
+
+  def testUnsortedSegmentSum2DIndices3DData(self):
+    for dtype in self.numeric_types:
+      data = np.array(
+          [[[0, 1, 2], [10, 11, 12]], [[100, 101, 102], [110, 111, 112]], [[
+              80, 81, 82
+          ], [123, 124, 125]], [[103, 104, 105], [106, 107, 108]]],
+          dtype=dtype)
+      indices = np.array([[3, 5], [3, 1], [5, 0], [6, 2]], dtype=np.int32)
+      num_segments = 8
+      y = self._unsortedSegmentSum(data, indices, num_segments)
+      self.assertAllClose(
+          np.array(
+              [[123, 124, 125], [110, 111, 112], [106, 107, 108], [
+                  100, 102, 104
+              ], [0, 0, 0.], [90, 92, 94], [103, 104, 105], [0, 0, 0]],
+              dtype=dtype), y)
+
+  def testUnsortedSegmentSum1DIndices3DData(self):
+    for dtype in self.numeric_types:
+      data = np.array(
+          [[[0, 1, 2], [10, 11, 12]], [[100, 101, 102], [110, 111, 112]], [[
+              120, 121, 122
+          ], [123, 124, 125]], [[103, 104, 105], [106, 107, 108]]],
+          dtype=dtype)
+      indices = np.array([3, 0, 2, 5], dtype=np.int32)
+      num_segments = 6
+      y = self._unsortedSegmentSum(data, indices, num_segments)
+      self.assertAllClose(
+          np.array(
+              [[[100, 101, 102.], [110, 111, 112]], [[0, 0, 0], [0, 0, 0]],
+               [[120, 121, 122], [123, 124, 125]], [[0, 1, 2.], [10, 11, 12]],
+               [[0, 0, 0], [0, 0, 0]], [[103, 104, 105], [106, 107, 108]]],
+              dtype=dtype), y)
+
+  def testUnsortedSegmentSumShapeError(self):
+    for dtype in self.numeric_types:
+      data = np.ones((4, 8, 7), dtype=dtype)
+      indices = np.ones((3, 2), dtype=np.int32)
+      num_segments = 4
+      self.assertRaises(
+          ValueError,
+          functools.partial(self._segmentReduction,
+                            math_ops.unsorted_segment_sum, data, indices,
+                            num_segments))
+
+  def testUnsortedSegmentOps1DIndices1DDataNegativeIndices(self):
+    """Tests for min, max, and prod ops.
+
+    These share most of their implementation with sum, so we only test basic
+    functionality.
+    """
+    for dtype in self.numeric_types:
+      self.assertAllClose(
+          np.array([8, 3, 1, 0], dtype=dtype),
+          self._unsortedSegmentProd(
+              np.array([0, 1, 2, 3, 4, 5, 6], dtype=dtype),
+              np.array([3, -1, 0, 1, 0, -1, 3], dtype=np.int32), 4))
+
+    for dtype in self.int_types | self.float_types:
+      minval = dtypes.as_dtype(dtype).min
+      maxval = dtypes.as_dtype(dtype).max
+
+      self.assertAllClose(
+          np.array([2, 3, maxval, 0], dtype=dtype),
+          self._unsortedSegmentMin(
+              np.array([0, 1, 2, 3, 4, 5, 6], dtype=dtype),
+              np.array([3, -1, 0, 1, 0, -1, 3], dtype=np.int32), 4))
+      self.assertAllClose(
+          np.array([4, 3, minval, 6], dtype=dtype),
+          self._unsortedSegmentMax(
+              np.array([0, 1, 2, 3, 4, 5, 6], dtype=dtype),
+              np.array([3, -1, 0, 1, 0, -1, 3], dtype=np.int32), 4))
+
+  def testNaNHandling(self):
+    for dtype in self.float_types:
+      hi = dtypes.as_dtype(dtype).max
+      lo = dtypes.as_dtype(dtype).min
+      if self.device == "TPU" and dtype == np.float64:
+        hi = np.inf
+        lo = -np.inf
+      self.assertAllClose(
+          np.array([1.0], dtype=dtype),
+          self._unsortedSegmentMin(
+              np.array([1.0, np.nan, 3.0], dtype=dtype),
+              np.array([0, 0, 0], dtype=np.int32), 1))
+      self.assertAllClose(
+          np.array([3.0], dtype=dtype),
+          self._unsortedSegmentMax(
+              np.array([np.nan, 1.0, 3.0], dtype=dtype),
+              np.array([0, 0, 0], dtype=np.int32), 1))
+      self.assertAllClose(
+          np.array([hi], dtype=dtype),
+          self._unsortedSegmentMin(
+              np.array([np.nan, np.nan], dtype=dtype),
+              np.array([0, 0], dtype=np.int32), 1))
+      self.assertAllClose(
+          np.array([lo], dtype=dtype),
+          self._unsortedSegmentMax(
+              np.array([np.nan, np.nan], dtype=dtype),
+              np.array([0, 0], dtype=np.int32), 1))
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/self_adjoint_eig_op_test.py

@@ -0,0 +1,59 @@
+# 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 tensorflow.ops.self_adjoint_eig."""
+
+import itertools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import linalg_ops
+from tensorflow.python.platform import test
+
+
+class SelfAdjointEigOpTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _test(self, dtype, shape):
+    np.random.seed(1)
+    x_np = np.random.uniform(
+        low=-1.0, high=1.0, size=np.prod(shape)).reshape(shape).astype(dtype)
+    x_np = x_np + np.swapaxes(x_np, -1, -2)
+    n = shape[-1]
+
+    e_np, _ = np.linalg.eigh(x_np)
+    with self.session() as sess:
+      x_tf = array_ops.placeholder(dtype)
+      with self.test_scope():
+        e, v = linalg_ops.self_adjoint_eig(x_tf)
+      e_val, v_val = sess.run([e, v], feed_dict={x_tf: x_np})
+
+      v_diff = np.matmul(v_val, np.swapaxes(v_val, -1, -2)) - np.eye(n)
+      self.assertAlmostEqual(np.mean(v_diff**2), 0.0, delta=1e-6)
+      self.assertAlmostEqual(np.mean((e_val - e_np)**2), 0.0, delta=1e-6)
+
+  SIZES = [1, 2, 5, 10, 32]
+  DTYPES = [np.float32]
+  PARAMS = itertools.product(SIZES, DTYPES)
+
+  @parameterized.parameters(*PARAMS)
+  def testSelfAdjointEig(self, n, dtype):
+    for batch_dims in [(), (3,)] + [(3, 2)] * (n < 10):
+      self._test(dtype, batch_dims + (n, n))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/sharding_util_ops_test.py

@@ -0,0 +1,713 @@
+# Copyright 2021 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 sharding util ops (XlaSplitND, XlaConcatND)."""
+
+from typing import Any, List, Optional
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.client.session import Session
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework.ops import control_dependencies
+from tensorflow.python.framework.tensor import Tensor
+from tensorflow.python.ops import gen_tpu_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def create_tensor_split_graph(
+    sess: Session,
+    input_value: Any,
+    input_dtype: Any,
+    num_outputs: int,
+    num_splits: List[int],
+    paddings: Optional[List[int]] = None) -> List[Tensor]:
+  del sess
+
+  const_input_op = constant_op.constant(input_value, dtype=input_dtype)
+  return gen_tpu_ops.xla_split_nd(
+      const_input_op, num_outputs, num_splits, paddings=paddings)
+
+
+def create_resource_split_graph(
+    sess: Session,
+    input_value: Any,
+    input_dtype: Any,
+    num_outputs: int,
+    num_splits: List[int],
+    paddings: Optional[List[int]] = None) -> List[Tensor]:
+  variable = resource_variable_ops.ResourceVariable(
+      initial_value=input_value, dtype=input_dtype)
+  sess.run(variables.variables_initializer([variable]))
+  return gen_tpu_ops.read_variable_xla_split_nd(
+      variable.handle, input_dtype, num_outputs, num_splits, paddings=paddings)
+
+
+class XlaSplitNDOpTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testSplitDimensionZero(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[[0]]],
+            input_dtype=dtype,
+            num_outputs=1,
+            num_splits=[1, 1, 0])
+        with self.assertRaisesOpError('index 2 must be positive, but got 0'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testSplitDimensionNegative(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[[0]]],
+            input_dtype=dtype,
+            num_outputs=1,
+            num_splits=[1, -1, 1])
+        with self.assertRaisesOpError('index 1 must be positive, but got -1'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testNumOutputsMismatch(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[0, 1],
+            input_dtype=dtype,
+            num_outputs=1,
+            num_splits=[2])
+        with self.assertRaisesOpError('\'N\' must match number of slices 2'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testPaddingsLengthMismatch(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[0, 1], [2, 3]],
+            input_dtype=dtype,
+            num_outputs=4,
+            num_splits=[2, 2],
+            paddings=[0])
+        with self.assertRaisesOpError('length 2, but got 1'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testPaddingsNegative(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[0, 1], [2, 3]],
+            input_dtype=dtype,
+            num_outputs=4,
+            num_splits=[2, 2],
+            paddings=[0, -1])
+        with self.assertRaisesOpError('non-negative, but got -1 at index 1'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testInputRankSplitMismatch(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[0, 1], [2, 3]],
+            input_dtype=dtype,
+            num_outputs=8,
+            num_splits=[2, 2, 2])
+        with self.assertRaisesOpError(
+            '\'num_splits\' length 3, but got rank 2'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testDimNotEvenlySplit(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[0, 1], [2, 3], [4, 5], [6, 7]],
+            input_dtype=dtype,
+            num_outputs=6,
+            num_splits=[3, 2])
+        with self.assertRaisesOpError('divisible by \'num_splits\' 3'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testDimWithPaddingNotEvenlySplit(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[0, 1], [2, 3], [4, 5], [6, 7]],
+            input_dtype=dtype,
+            num_outputs=4,
+            num_splits=[2, 2],
+            paddings=[0, 1])
+        with self.assertRaisesOpError('divisible by \'num_splits\' 2'):
+          sess.run(split)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testNoSplits(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[[0, 1], [2, 3]], [[4, 5], [6, 7]]],
+            input_dtype=dtype,
+            num_outputs=1,
+            num_splits=[1, 1, 1])
+        results = sess.run(split)
+      self.assertLen(results, 1)
+      self.assertAllClose(results[0], [[[0, 1], [2, 3]], [[4, 5], [6, 7]]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testNoSplitsWithPadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[[0]], [[1]]],
+            input_dtype=dtype,
+            num_outputs=1,
+            num_splits=[1, 1, 1],
+            paddings=[0, 1, 1])
+        results = sess.run(split)
+      self.assertLen(results, 1)
+      self.assertAllClose(results[0], [[[0, 0], [0, 0]], [[1, 0], [0, 0]]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testSplitNoPadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[
+                [0, 1, 2, 3],
+                [4, 5, 6, 7],
+                [8, 9, 10, 11],
+                [12, 13, 14, 15],
+            ],
+            input_dtype=dtype,
+            num_outputs=4,
+            num_splits=[2, 2])
+        results = sess.run(split)
+      self.assertLen(results, 4)
+      self.assertAllClose(results[0], [[0, 1], [4, 5]])
+      self.assertAllClose(results[1], [[2, 3], [6, 7]])
+      self.assertAllClose(results[2], [[8, 9], [12, 13]])
+      self.assertAllClose(results[3], [[10, 11], [14, 15]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testSplitPartialPadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[
+                [0, 1, 2],
+                [3, 4, 5],
+                [6, 7, 8],
+            ],
+            input_dtype=dtype,
+            num_outputs=4,
+            num_splits=[2, 2],
+            paddings=[1, 1])
+        results = sess.run(split)
+      self.assertLen(results, 4)
+      self.assertAllClose(results[0], [[0, 1], [3, 4]])
+      self.assertAllClose(results[1], [[2, 0], [5, 0]])
+      self.assertAllClose(results[2], [[6, 7], [0, 0]])
+      self.assertAllClose(results[3], [[8, 0], [0, 0]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_split_graph),
+                                  ('Resource', create_resource_split_graph))
+  def testSplitCompletePadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=[[0], [1]],
+            input_dtype=dtype,
+            num_outputs=4,
+            num_splits=[2, 2],
+            paddings=[2, 3])
+        results = sess.run(split)
+      self.assertLen(results, 4)
+      self.assertAllClose(results[0], [[0, 0], [1, 0]])
+      self.assertAllClose(results[1], [[0, 0], [0, 0]])
+      self.assertAllClose(results[2], [[0, 0], [0, 0]])
+      self.assertAllClose(results[3], [[0, 0], [0, 0]])
+
+  @parameterized.named_parameters(
+      ('1Tensor', create_tensor_split_graph, 1),
+      ('2Tensor', create_tensor_split_graph, 2),
+      ('3Tensor', create_tensor_split_graph, 3),
+      ('4Tensor', create_tensor_split_graph, 4),
+      ('5Tensor', create_tensor_split_graph, 5),
+      ('6Tensor', create_tensor_split_graph, 6),
+      ('7Tensor', create_tensor_split_graph, 7),
+      ('8Tensor', create_tensor_split_graph, 8),
+      ('1Resource', create_resource_split_graph, 1),
+      ('2Resource', create_resource_split_graph, 2),
+      ('3Resource', create_resource_split_graph, 3),
+      ('4Resource', create_resource_split_graph, 4),
+      ('5Resource', create_resource_split_graph, 5),
+      ('6Resource', create_resource_split_graph, 6),
+      ('7Resource', create_resource_split_graph, 7),
+      ('8Resource', create_resource_split_graph, 8),
+  )
+  def testRanked(self, graph_fn, rank):
+    num_splits = [2] * rank
+    num_outputs = 2 << (rank - 1)
+    input_value = np.reshape(np.arange(np.prod(num_splits)), num_splits)
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        split = graph_fn(
+            sess,
+            input_value=input_value,
+            input_dtype=dtype,
+            num_outputs=num_outputs,
+            num_splits=num_splits)
+        results = sess.run(split)
+      self.assertLen(results, num_outputs)
+      for i, result in enumerate(results):
+        expected_output = np.reshape(i, [1] * rank).astype(dtype)
+        self.assertAllClose(result, expected_output)
+
+
+def create_tensor_concat_graph(
+    sess: Session,
+    input_values: List[Any],
+    input_dtype: Any,
+    num_concats: List[int],
+    paddings: Optional[List[int]] = None,
+    output_shape: Optional[List[int]] = None) -> Tensor:
+  del sess
+  del output_shape
+
+  const_input_ops = [
+      constant_op.constant(i, dtype=input_dtype) for i in input_values
+  ]
+  return gen_tpu_ops.xla_concat_nd(const_input_ops, num_concats, paddings)
+
+
+def create_resource_concat_graph(
+    sess: Session,
+    input_values: List[Any],
+    input_dtype: Any,
+    num_concats: List[int],
+    paddings: Optional[List[int]] = None,
+    output_shape: Optional[List[int]] = None) -> Tensor:
+  variable_shape = [] if output_shape is None else output_shape
+  variable = resource_variable_ops.ResourceVariable(
+      initial_value=np.zeros(variable_shape, dtype=input_dtype),
+      dtype=input_dtype)
+  sess.run(variables.variables_initializer([variable]))
+  const_input_ops = [
+      constant_op.constant(i, dtype=input_dtype) for i in input_values
+  ]
+  concat = gen_tpu_ops.assign_variable_xla_concat_nd(variable.handle,
+                                                     const_input_ops,
+                                                     num_concats, paddings)
+  with control_dependencies([concat]):
+    return variable.read_value()
+
+
+class XlaConcatNDOpTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testConcatDimensionZero(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[[[0]]]],
+            input_dtype=dtype,
+            num_concats=[1, 1, 0])
+        with self.assertRaisesOpError('index 2 must be positive, but got 0'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testConcatDimensionNegative(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[[[0]]]],
+            input_dtype=dtype,
+            num_concats=[1, -1, 1])
+        with self.assertRaisesOpError('index 1 must be positive, but got -1'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testNumInputsMismatch(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess, input_values=[[0, 1]], input_dtype=dtype, num_concats=[2])
+        with self.assertRaisesOpError('\'N\' must match number of slices 2'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testPaddingsLengthMismatch(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[[0, 1], [2, 3]]],
+            input_dtype=dtype,
+            num_concats=[1, 1],
+            paddings=[0])
+        with self.assertRaisesOpError('length 2, but got 1'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testPaddingsNegative(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[[0, 1], [2, 3]]],
+            input_dtype=dtype,
+            num_concats=[1, 1],
+            paddings=[0, -1])
+        with self.assertRaisesOpError('non-negative, but got -1 at index 1'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testInputRankConcatMismatch(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess, input_values=[[0]], input_dtype=dtype, num_concats=[1, 1])
+        with self.assertRaisesOpError(
+            '\'num_concats\' length 2, but got rank 1'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testDifferentShapedInputs(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[0], [1, 2]],
+            input_dtype=dtype,
+            num_concats=[2])
+        with self.assertRaisesOpError(
+            r'same expected shape \[1\], but got \[2\] at index 1'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testPaddingExceedsOutputDimSize(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[0]],
+            input_dtype=dtype,
+            num_concats=[1],
+            paddings=[2])
+        with self.assertRaisesOpError(
+            'exceed expected output shape dimension 1 at index 0, but got 2'):
+          sess.run(concat)
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testNoConcats(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]],
+            input_dtype=dtype,
+            num_concats=[1, 1, 1],
+            output_shape=[2, 2, 2])
+        result = sess.run(concat)
+      self.assertAllClose(result, [[[0, 1], [2, 3]], [[4, 5], [6, 7]]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testNoConcatsWithPadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[[[[0, 1], [2, 3]], [[4, 5], [6, 7]]]],
+            input_dtype=dtype,
+            num_concats=[1, 1, 1],
+            output_shape=[1, 1, 1],
+            paddings=[1, 1, 1])
+        result = sess.run(concat)
+      self.assertAllClose(result, [[[0]]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testConcatNoPadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[
+                [[0, 1], [2, 3]],
+                [[4, 5], [6, 7]],
+                [[8, 9], [10, 11]],
+                [[12, 13], [14, 15]],
+            ],
+            input_dtype=dtype,
+            num_concats=[2, 2],
+            output_shape=[4, 4])
+        result = sess.run(concat)
+      self.assertAllClose(
+          result,
+          [[0, 1, 4, 5], [2, 3, 6, 7], [8, 9, 12, 13], [10, 11, 14, 15]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testConcatPartialPadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[
+                [[0, 1], [2, 3]],
+                [[4, 5], [6, 7]],
+                [[8, 9], [10, 11]],
+                [[12, 13], [14, 15]],
+            ],
+            input_dtype=dtype,
+            num_concats=[2, 2],
+            output_shape=[3, 3],
+            paddings=[1, 1])
+        result = sess.run(concat)
+      self.assertAllClose(result, [[0, 1, 4], [2, 3, 6], [8, 9, 12]])
+
+  @parameterized.named_parameters(('Tensor', create_tensor_concat_graph),
+                                  ('Resource', create_resource_concat_graph))
+  def testConcatCompletePadding(self, graph_fn):
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=[
+                [[0, 1], [2, 3]],
+                [[4, 5], [6, 7]],
+                [[8, 9], [10, 11]],
+                [[12, 13], [14, 15]],
+            ],
+            input_dtype=dtype,
+            num_concats=[2, 2],
+            output_shape=[2, 2],
+            paddings=[2, 2])
+        result = sess.run(concat)
+      self.assertAllClose(result, [[0, 1], [2, 3]])
+
+  @parameterized.named_parameters(
+      ('1Tensor', create_tensor_concat_graph, 1),
+      ('2Tensor', create_tensor_concat_graph, 2),
+      ('3Tensor', create_tensor_concat_graph, 3),
+      ('4Tensor', create_tensor_concat_graph, 4),
+      ('5Tensor', create_tensor_concat_graph, 5),
+      ('6Tensor', create_tensor_concat_graph, 6),
+      ('7Tensor', create_tensor_concat_graph, 7),
+      ('8Tensor', create_tensor_concat_graph, 8),
+      ('1Resource', create_resource_concat_graph, 1),
+      ('2Resource', create_resource_concat_graph, 2),
+      ('3Resource', create_resource_concat_graph, 3),
+      ('4Resource', create_resource_concat_graph, 4),
+      ('5Resource', create_resource_concat_graph, 5),
+      ('6Resource', create_resource_concat_graph, 6),
+      ('7Resource', create_resource_concat_graph, 7),
+      ('8Resource', create_resource_concat_graph, 8),
+  )
+  def testRanked(self, graph_fn, rank):
+    num_concats = [2] * rank
+    num_inputs = 2 << (rank - 1)
+    input_values = [np.reshape(i, [1] * rank) for i in range(num_inputs)]
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        concat = graph_fn(
+            sess,
+            input_values=input_values,
+            input_dtype=dtype,
+            num_concats=num_concats,
+            output_shape=num_concats)
+        result = sess.run(concat)
+      expected_output = np.arange(0,
+                                  num_inputs).reshape(num_concats).astype(dtype)
+      self.assertAllClose(result, expected_output)
+
+
+def create_tensor_roundtrip_graph(
+    sess: Session,
+    value: Any,
+    dtype: Any,
+    num_partitions: List[int],
+    paddings: Optional[List[int]] = None) -> Tensor:
+  del sess
+
+  const_input_op = constant_op.constant(value, dtype=dtype)
+  split = gen_tpu_ops.xla_split_nd(
+      const_input_op,
+      np.prod(num_partitions),
+      num_partitions,
+      paddings=paddings)
+  concat = gen_tpu_ops.xla_concat_nd(split, num_partitions, paddings)
+  return math_ops.equal(const_input_op, concat)
+
+
+def create_resource_roundtrip_graph(
+    sess: Session,
+    value: Any,
+    dtype: Any,
+    num_partitions: List[int],
+    paddings: Optional[List[int]] = None) -> Tensor:
+  variable = resource_variable_ops.ResourceVariable(
+      initial_value=value, dtype=dtype)
+  sess.run(variables.variables_initializer([variable]))
+  split = gen_tpu_ops.read_variable_xla_split_nd(
+      variable.handle,
+      dtype,
+      np.prod(num_partitions),
+      num_partitions,
+      paddings=paddings)
+  concat = gen_tpu_ops.assign_variable_xla_concat_nd(variable.handle, split,
+                                                     num_partitions, paddings)
+  with control_dependencies([concat]):
+    return math_ops.equal(variable.read_value(),
+                          constant_op.constant(value, dtype=dtype))
+
+
+class XlaSplitConcatNDTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('1Tensor', create_tensor_roundtrip_graph, 1),
+      ('2Tensor', create_tensor_roundtrip_graph, 2),
+      ('3Tensor', create_tensor_roundtrip_graph, 3),
+      ('4Tensor', create_tensor_roundtrip_graph, 4),
+      ('5Tensor', create_tensor_roundtrip_graph, 5),
+      ('6Tensor', create_tensor_roundtrip_graph, 6),
+      ('7Tensor', create_tensor_roundtrip_graph, 7),
+      ('8Tensor', create_tensor_roundtrip_graph, 8),
+      ('1Resource', create_resource_roundtrip_graph, 1),
+      ('2Resource', create_resource_roundtrip_graph, 2),
+      ('3Resource', create_resource_roundtrip_graph, 3),
+      ('4Resource', create_resource_roundtrip_graph, 4),
+      ('5Resource', create_resource_roundtrip_graph, 5),
+      ('6Resource', create_resource_roundtrip_graph, 6),
+      ('7Resource', create_resource_roundtrip_graph, 7),
+      ('8Resource', create_resource_roundtrip_graph, 8),
+  )
+  def testNoPadding(self, graph_fn, rank):
+    num_partitions = [2] * rank
+    shape = [4] * rank
+    value = np.arange(0, np.prod(shape)).reshape(shape)
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        validate = graph_fn(sess, value, dtype, num_partitions)
+        result = sess.run(validate)
+      self.assertAllEqual(result, np.broadcast_to(True, shape))
+
+  @parameterized.named_parameters(
+      ('1Tensor', create_tensor_roundtrip_graph, 1),
+      ('2Tensor', create_tensor_roundtrip_graph, 2),
+      ('3Tensor', create_tensor_roundtrip_graph, 3),
+      ('4Tensor', create_tensor_roundtrip_graph, 4),
+      ('5Tensor', create_tensor_roundtrip_graph, 5),
+      ('6Tensor', create_tensor_roundtrip_graph, 6),
+      ('7Tensor', create_tensor_roundtrip_graph, 7),
+      ('8Tensor', create_tensor_roundtrip_graph, 8),
+      ('1Resource', create_resource_roundtrip_graph, 1),
+      ('2Resource', create_resource_roundtrip_graph, 2),
+      ('3Resource', create_resource_roundtrip_graph, 3),
+      ('4Resource', create_resource_roundtrip_graph, 4),
+      ('5Resource', create_resource_roundtrip_graph, 5),
+      ('6Resource', create_resource_roundtrip_graph, 6),
+      ('7Resource', create_resource_roundtrip_graph, 7),
+      ('8Resource', create_resource_roundtrip_graph, 8),
+  )
+  def testPartialPadding(self, graph_fn, rank):
+    num_partitions = [2] * rank
+    shape = [4] * rank
+    value = np.arange(0, np.prod(shape)).reshape(shape)
+    paddings = [2] * rank
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        validate = graph_fn(sess, value, dtype, num_partitions, paddings)
+        result = sess.run(validate)
+      self.assertAllEqual(result, np.broadcast_to(True, shape))
+
+  @parameterized.named_parameters(
+      ('1Tensor', create_tensor_roundtrip_graph, 1),
+      ('2Tensor', create_tensor_roundtrip_graph, 2),
+      ('3Tensor', create_tensor_roundtrip_graph, 3),
+      ('4Tensor', create_tensor_roundtrip_graph, 4),
+      ('5Tensor', create_tensor_roundtrip_graph, 5),
+      ('6Tensor', create_tensor_roundtrip_graph, 6),
+      ('7Tensor', create_tensor_roundtrip_graph, 7),
+      ('8Tensor', create_tensor_roundtrip_graph, 8),
+      ('1Resource', create_resource_roundtrip_graph, 1),
+      ('2Resource', create_resource_roundtrip_graph, 2),
+      ('3Resource', create_resource_roundtrip_graph, 3),
+      ('4Resource', create_resource_roundtrip_graph, 4),
+      ('5Resource', create_resource_roundtrip_graph, 5),
+      ('6Resource', create_resource_roundtrip_graph, 6),
+      ('7Resource', create_resource_roundtrip_graph, 7),
+      ('8Resource', create_resource_roundtrip_graph, 8),
+  )
+  def testCompletePadding(self, graph_fn, rank):
+    num_partitions = [2] * rank
+    shape = [4] * rank
+    value = np.arange(0, np.prod(shape)).reshape(shape)
+    paddings = [4] * rank
+    for dtype in self.numeric_types:
+      with self.session() as sess, self.device_scope():
+        validate = graph_fn(sess, value, dtype, num_partitions, paddings)
+        result = sess.run(validate)
+      self.assertAllEqual(result, np.broadcast_to(True, shape))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/compiler/tests/slice_ops_test.py

@@ -0,0 +1,302 @@
+# 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.
+# ==============================================================================
+"""Tests for slicing."""
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import googletest
+
+
+class SliceTest(xla_test.XLATestCase):
+
+  def test1D(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[10])
+        with self.test_scope():
+          o = array_ops.slice(i, [2], [4])
+        params = {
+            i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([2, 3, 4, 5], result)
+
+  def testZeroSlice(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[2])
+        with self.test_scope():
+          o = array_ops.slice(i, [0], [0])
+        params = {
+            i: [0, 1],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([], result)
+
+  def test3D(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[3, 3, 10])
+        with self.test_scope():
+          o = array_ops.slice(i, [1, 2, 2], [1, 1, 4])
+        params = {
+            i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+                 [9, 8, 7, 6, 5, 4, 3, 2, 1, 0],
+                 [5, 3, 1, 7, 9, 2, 4, 6, 8, 0]],
+                [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
+                 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+                 [8, 7, 6, 5, 4, 3, 2, 1, 8, 7]],
+                [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5],
+                 [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
+                 [9, 8, 7, 9, 8, 7, 9, 8, 7, 9]]]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[[6, 5, 4, 3]]], result)
+
+  def test3DWithDynamicBegin(self):
+    """Tests a slice where the start offset is not known at compile time."""
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[3, 3, 10])
+        begin = array_ops.placeholder(dtypes.int32, shape=[3])
+        with self.test_scope():
+          o = array_ops.slice(i, begin, [1, 1, 4])
+        params = {
+            i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+                 [9, 8, 7, 6, 5, 4, 3, 2, 1, 0],
+                 [5, 3, 1, 7, 9, 2, 4, 6, 8, 0]],
+                [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
+                 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+                 [8, 7, 6, 5, 4, 3, 2, 1, 8, 7]],
+                [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5],
+                 [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
+                 [9, 8, 7, 9, 8, 7, 9, 8, 7, 9]]],
+            begin: [1, 2, 2]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[[6, 5, 4, 3]]], result)
+
+  def test3DWithDynamicBeginAndNegativeSize(self):
+    """Tests a slice where `begin` is fed dynamically and `size` contains -1."""
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[3, 3, 10])
+        begin = array_ops.placeholder(dtypes.int32, shape=[3])
+        with self.test_scope():
+          o = array_ops.slice(i, begin, [1, -1, 4])
+        params = {
+            i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+                 [9, 8, 7, 6, 5, 4, 3, 2, 1, 0],
+                 [5, 3, 1, 7, 9, 2, 4, 6, 8, 0]],
+                [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
+                 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+                 [8, 7, 6, 5, 4, 3, 2, 1, 8, 7]],
+                [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5],
+                 [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
+                 [9, 8, 7, 9, 8, 7, 9, 8, 7, 9]]],
+            begin: [1, 1, 2]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[[1, 1, 1, 1], [6, 5, 4, 3]]], result)
+
+
+class StridedSliceTest(xla_test.XLATestCase):
+
+  def test1D(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[10])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [2], [6], [2])
+        params = {
+            i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([2, 4], result)
+
+  def test1DDynamic(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[10])
+        begin = array_ops.placeholder(dtypes.int32, shape=[1])
+        with self.test_scope():
+          end = math_ops.add(begin, [1])
+          o = array_ops.strided_slice(i, begin, end, [1])
+        params = {
+            i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            begin: [0]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([0], result)
+
+  def testDynamicSliceOutOfBounds(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[10])
+        begin = array_ops.placeholder(dtypes.int32, shape=[1])
+        end = array_ops.placeholder(dtypes.int32, shape=[1])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, begin, end, [1])
+        params = {i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], begin: [5], end: [15]}
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([5, 6, 7, 8, 9], result)
+
+  def test1DNegativeStride(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[10])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [6], [2], [-2])
+        params = {
+            i: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([6, 4], result)
+
+  def test2DDegenerate(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[2, 3])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [-1, 0], [0, 3])
+        params = {
+            i: [[0, 1, 2],
+                [3, 4, 5]]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertEqual(tensor_shape.TensorShape((0, 3)), result.shape)
+
+  def test2DDegenerateNegativeStride(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[2, 3])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [0, 0], [-1, 3], [-1, 1])
+        params = {
+            i: [[0, 1, 2],
+                [3, 4, 5]]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertEqual(tensor_shape.TensorShape((0, 3)), result.shape)
+
+  def test2DFullSlice(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        with self.test_scope():
+          i = array_ops.placeholder(dtype, shape=[2, 4])
+          begin = array_ops.placeholder(dtypes.int32, shape=[2])
+          end = math_ops.add(begin, [1, 1])
+          o = array_ops.strided_slice(i, begin, end, [1, 1])
+        params = {
+            i: [[0, 1, 2, 3], [4, 5, 6, 7]],
+            begin: [1, 1]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[5]], result)
+
+  def test3D(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[3, 3, 10])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [0, 2, 2], [2, 3, 6], [1, 1, 2])
+        params = {
+            i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+                 [9, 8, 7, 6, 5, 4, 3, 2, 1, 0],
+                 [5, 3, 1, 7, 9, 2, 4, 6, 8, 0]],
+                [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
+                 [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+                 [8, 7, 6, 5, 4, 3, 2, 1, 8, 7]],
+                [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5],
+                 [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
+                 [9, 8, 7, 9, 8, 7, 9, 8, 7, 9]]]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[[1, 9]], [[6, 4]]], result)
+
+  def test3DNegativeStride(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[3, 4, 10])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [2, 2, 6], [0, 0, 2], [-1, -1, -2])
+        params = {
+            i: [[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+                 [9, 8, 7, 6, 5, 4, 3, 2, 1, 0],
+                 [5, 3, 1, 7, 9, 2, 4, 6, 8, 0],
+                 [4, 5, 2, 4, 3, 7, 6, 8, 9, 4]],
+                [[5, 5, 5, 5, 5, 5, 5, 5, 5, 5],
+                 [4, 3, 4, 5, 7, 6, 5, 3, 4, 5],
+                 [8, 7, 6, 5, 4, 3, 2, 1, 8, 7],
+                 [7, 1, 7, 1, 8, 1, 8, 1, 3, 1]],
+                [[7, 5, 7, 5, 7, 5, 7, 5, 7, 5],
+                 [1, 2, 1, 2, 1, 2, 1, 2, 1, 2],
+                 [9, 8, 7, 9, 8, 7, 9, 8, 7, 9],
+                 [9, 9, 5, 5, 6, 6, 3, 3, 6, 6]]]
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([[[9, 8],
+                              [1, 1]],
+                             [[2, 4],
+                              [5, 7]]], result)
+
+  # Test shrink_axis_mask. This `strided_slice` call is equivalent to `i[1,:]`.
+  def testShrinkAxisMask(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[2, 3])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [1, 0], [10, 3], shrink_axis_mask=1)
+        params = {
+            i: [[0, 1, 2], [3, 4, 5]],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([3, 4, 5], result)
+
+  # Test shrink_axis_mask with the range for the second dimension implicit.
+  # This `strided_slice` call is equivalent to `i[1]`.
+  def testShrinkAxisMaskImplicitRange(self):
+    for dtype in self.numeric_types:
+      with self.session():
+        i = array_ops.placeholder(dtype, shape=[2, 3])
+        with self.test_scope():
+          o = array_ops.strided_slice(i, [1], [10], shrink_axis_mask=1)
+        params = {
+            i: [[0, 1, 2], [3, 4, 5]],
+        }
+        result = o.eval(feed_dict=params)
+
+        self.assertAllEqual([3, 4, 5], result)
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/compiler/tests/sort_ops_test.py

@@ -0,0 +1,482 @@
+# Copyright 2018 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 sorting operators."""
+
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.compiler.tf2xla.python import xla
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import sort_ops
+from tensorflow.python.platform import test
+
+ALL_KEY_TYPES = [
+    dtypes.bfloat16.as_numpy_dtype, np.float16, np.float32, np.float64,
+    np.int32, np.uint32, np.int16, np.uint16, np.int8, np.uint8
+]
+
+
+class XlaSortOpTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def _assertOpOutputMatchesExpected(self, op, args, expected):
+    """Tests that op(*args) == expected."""
+    with self.session() as session:
+      with self.test_scope():
+        placeholders = [
+            array_ops.placeholder(dtypes.as_dtype(arg.dtype), arg.shape)
+            for arg in args
+        ]
+        feeds = {placeholders[i]: args[i] for i in range(0, len(args))}
+        output = op(*placeholders)
+        if isinstance(output, tensor.Tensor):
+          output = [output]
+
+      results = session.run(output, feeds)
+      for result, v in zip(results, expected):
+        self.assertAllClose(v, result, rtol=1e-3)
+
+  def _shuffled_arange(self, shape, dtype):
+    x = np.arange(np.prod(shape), dtype=dtype)
+    np.random.shuffle(x)
+    return x.reshape(shape)
+
+  def _supported_key_types(self):
+    supported_key_types = set(ALL_KEY_TYPES)
+    res = supported_key_types.intersection(self.numeric_types)
+    assert res
+    return res
+
+  def testSort(self):
+    for dtype in self._supported_key_types():
+      x = self._shuffled_arange((101,), dtype)
+      self._assertOpOutputMatchesExpected(
+          xla.sort, [x], expected=[np.arange(101, dtype=dtype)])
+
+  def testKeyValueSort(self):
+    for key_type in self._supported_key_types():
+      for value_type in self._supported_key_types():
+        if key_type == np.uint8 or value_type == np.uint8:
+          # I do not understand why the test fails on uint8. We plan to
+          # deprecate xla.key_value_sort in favor of xla.variadic_sort anyway.
+          continue
+        x = self._shuffled_arange((101,), key_type)
+        y = (-x).astype(value_type)
+        self._assertOpOutputMatchesExpected(
+            xla.key_value_sort, [x, y],
+            expected=[
+                np.arange(101, dtype=key_type),
+                -np.arange(101, dtype=value_type)
+            ])
+
+  # Flip is the only reliable way to get a descending sort across any dimension.
+  # 1. -np.sort(-x) doesn't work with unsigned integers.
+  # 2. np.sort(x, axis=a)[::-1] is not generic over axis wher reversing array
+  # 3. x.argsort() either requires "-1" (first option) or flip, so the same.
+  def _descendingSort(self, x, dimension):
+    b = np.sort(x, axis=dimension)
+    return np.flip(b, axis=dimension)
+
+  @parameterized.parameters(0, 1, 2)
+  def testMisleadingComparator(self, dimension):
+    shape = (4, 3, 4)
+    for key_type in self._supported_key_types():
+      x = self._shuffled_arange(shape, key_type)
+      expected = self._descendingSort(x, dimension)
+
+      # pylint: disable=cell-var-from-loop
+      @function.Defun(key_type, key_type)
+      def compare_gt(x1, x2):
+        return x2 < x1  # "greater than" with misleading "<" sign
+
+      def wrap_sort(x):
+        return xla.variadic_sort([x],
+                                 dimension=dimension,
+                                 is_stable=False,
+                                 comparator=compare_gt)
+
+      self._assertOpOutputMatchesExpected(wrap_sort, [x], expected=[expected])
+
+  @parameterized.parameters(0, 1, 2)
+  def testVariadicSortDimension(self, dimension):
+    shape = (2, 3, 4)
+    for key_type in self._supported_key_types():
+      x = self._shuffled_arange(shape, key_type)
+      expected = np.sort(x, axis=dimension)
+
+      @function.Defun(key_type, key_type)
+      def compare_lt(x1, x2):
+        return x1 < x2
+
+      def wrap_sort(x):
+        return xla.variadic_sort([x],
+                                 dimension=dimension,
+                                 is_stable=False,
+                                 comparator=compare_lt)
+
+      self._assertOpOutputMatchesExpected(wrap_sort, [x], expected=[expected])
+
+  @parameterized.parameters(0, 1, 2)
+  def testVariadicSortReverse(self, dimension):
+    shape = (100, 3, 4)
+    for key_type in self._supported_key_types():
+      x = self._shuffled_arange(shape, key_type)
+      expected = self._descendingSort(x, dimension)
+
+      @function.Defun(key_type, key_type)
+      def compare_gt(x1, x2):
+        return x1 > x2
+
+      def wrap_sort(x):
+        return xla.variadic_sort([x],
+                                 dimension=dimension,
+                                 is_stable=False,
+                                 comparator=compare_gt)
+
+      self._assertOpOutputMatchesExpected(wrap_sort, [x], expected=[expected])
+
+  @parameterized.product(dimension=[0, 1, 2], key_type=ALL_KEY_TYPES)
+  def testVariadicSortSeveral(self, dimension, key_type):
+    if np.__version__ < "1.15":
+      raise unittest.SkipTest("np.take_along_axis was added in 1.15")
+    if key_type not in self._supported_key_types():
+      return
+    shape = (2, 3, 4)
+    for value_type_1 in self._supported_key_types():
+      for value_type_2 in self._supported_key_types():
+        inputs = [
+            self._shuffled_arange(shape, key_type),
+            self._shuffled_arange(shape, value_type_1),
+            self._shuffled_arange(shape, value_type_2)
+        ]
+
+        # The first array is sorted, and the others are shuffled the same way
+        sorted_indices = np.argsort(inputs[0], axis=dimension)
+        expected = [
+            np.take_along_axis(inp, sorted_indices, axis=dimension)
+            for inp in inputs
+        ]
+        self.assertAllEqual(np.sort(inputs[0], axis=dimension), expected[0])
+
+        @function.Defun(key_type, key_type, value_type_1, value_type_1,
+                        value_type_2, value_type_2)
+        def compare_lt(x1, x2, y1, y2, z1, z2):
+          del y1, y2, z1, z2
+          return x1 < x2
+
+        def wrap_sort(*args):
+          return xla.variadic_sort(
+              args,  # Pass the arguments as a tuple
+              comparator=compare_lt,
+              dimension=dimension,
+              is_stable=False)
+
+        self._assertOpOutputMatchesExpected(
+            wrap_sort, inputs, expected=expected)
+
+  @parameterized.parameters(ALL_KEY_TYPES)
+  @test_util.disable_mlir_bridge("Not supported yet")
+  def testVariadicSortLexicographic(self, key_type_2):
+    # Three inputs: the first two are used for lexicographic sort, and the
+    # third is just swapped accordingly.
+    # The first array will contain only 0 and 1, to test lexicographic order
+    if np.__version__ < "1.15":
+      raise unittest.SkipTest("np.take_along_axis was added in 1.15")
+    shape = (20,)
+    if key_type_2 not in self._supported_key_types():
+      return
+    for key_type_1 in [np.int16, np.uint16, np.int32, np.uint32]:
+      for value_type in self._supported_key_types():
+        inputs = [
+            # Ensure that some keys in the first input are equal
+            np.random.uniform(0, 2, shape).astype(key_type_1),
+            self._shuffled_arange(shape, key_type_2),
+            self._shuffled_arange(shape, value_type)
+        ]
+        # The first two arrays are sorted lexicographically, and the third
+        # is shuffled the same way
+        sorted_indices = np.argsort(100 * inputs[0] + inputs[1])
+        expected = [
+            np.take_along_axis(inp, sorted_indices, axis=0) for inp in inputs
+        ]
+
+        @function.Defun(key_type_1, key_type_1, key_type_2, key_type_2,
+                        value_type, value_type)
+        def compare_lexicographic(x1, x2, y1, y2, z1, z2):
+          del z1, z2
+          return math_ops.logical_or(
+              x1 < x2, math_ops.logical_and(math_ops.equal(x1, x2), y1 < y2))
+
+        def wrap_sort(*args):
+          return xla.variadic_sort(
+              args,  # Pass the arguments as a tuple
+              comparator=compare_lexicographic,
+              dimension=0,
+              is_stable=False)
+
+        self._assertOpOutputMatchesExpected(
+            wrap_sort, inputs, expected=expected)
+
+  @parameterized.product(dimension=[0, 1, 2], key_type=ALL_KEY_TYPES)
+  def testVariadicSortSeveralStable(self, dimension, key_type):
+    shape = (2, 3, 4)
+    if key_type not in self._supported_key_types():
+      return
+    for value_type_1 in self._supported_key_types():
+      for value_type_2 in self._supported_key_types():
+        # The first input is all 0s, there should be no changes for
+        # stable sort.
+        inputs = [
+            np.zeros(shape, key_type),
+            self._shuffled_arange(shape, value_type_1),
+            self._shuffled_arange(shape, value_type_2)
+        ]
+
+        @function.Defun(key_type, key_type, value_type_1, value_type_1,
+                        value_type_2, value_type_2)
+        def compare_lt(x1, x2, y1, y2, z1, z2):
+          del y1, y2, z1, z2
+          return x1 < x2
+
+        def wrap_sort(*args):
+          return xla.variadic_sort(
+              args,  # Pass the arguments as a tuple
+              comparator=compare_lt,
+              dimension=dimension,
+              is_stable=True)
+
+        self._assertOpOutputMatchesExpected(wrap_sort, inputs, expected=inputs)
+
+  @parameterized.product(dimension=[0, 1, 2], dtype=ALL_KEY_TYPES)
+  def testArgsort(self, dimension, dtype):
+    shape = (2, 3, 4)
+    if dtype not in self._supported_key_types():
+      return
+
+    def argsort(v, axis=dimension):
+      return sort_ops.argsort(v, axis, stable=True)
+
+    x = self._shuffled_arange(shape, dtype)
+    self._assertOpOutputMatchesExpected(
+        argsort, [x], expected=[np.argsort(x, axis=dimension, kind="stable")]
+    )
+
+  @parameterized.product(
+      dtype=[
+          dtypes.bfloat16.as_numpy_dtype,
+          np.float16,
+          np.float32,
+          np.float64,
+          np.int32,
+          np.uint32,
+          np.int64,
+          np.uint64,
+          np.uint8,
+          np.int8,
+      ],
+      rank=[1, 2, 3],
+  )
+  def testTopK(self, dtype, rank):
+    if dtype in self.numeric_types:
+      # Use small input size for bfloat16. Otherwise, we'll get duplicate values
+      # after conversion to bfloat16, so the possible resulting index array is
+      # no longer unique.
+      if dtype in (dtypes.bfloat16.as_numpy_dtype, np.float16):
+        array_size = 20
+        k_options = [0, 1, 2, 10, 20]
+      elif dtype in (dtypes.uint8.as_numpy_dtype, dtypes.int8.as_numpy_dtype):
+        array_size = 111
+        k_options = [0, 1, 2, 10, 20]
+      else:
+        array_size = 200 * 1000
+        k_options = [0, 1, 2, 10, 20, 100, 1000, 200 * 1000]
+
+      # Tile array to tensor of specified rank, then shuffle along the last dim
+      x = np.arange(array_size)
+      x = np.tile(x, (2,) * (rank - 1) + (1,))
+      np.apply_along_axis(np.random.shuffle, -1, x)
+
+      sorted_indices = x.argsort(axis=-1)[..., ::-1]
+      sorted_values = np.sort(x, axis=-1)[..., ::-1]
+      for k in k_options:
+        indices = sorted_indices[..., :k]
+        expected = sorted_values[..., :k]
+
+        def topk(v, k=k):
+          return nn_ops.top_k(v, k=k, sorted=True)
+
+        self._assertOpOutputMatchesExpected(
+            topk,
+            [x.astype(dtype)],
+            expected=[expected.astype(dtype), indices],
+        )
+
+  def testTopKZeros(self):
+    """Tests that positive and negative zeros sort correctly."""
+    supported_types = set(
+        [dtypes.bfloat16.as_numpy_dtype, np.float16, np.float32, np.float64])
+    for dtype in supported_types.intersection(self.numeric_types):
+      with self.session() as sess:
+        p = array_ops.placeholder(dtype)
+        with self.test_scope():
+          topk = nn_ops.top_k(p, k=4)
+        results = sess.run(
+            topk,
+            {p: np.array([0., -0., 0., 3., -0., -4., 0., -0.], dtype=dtype)})
+        self.assertAllEqual(np.array([3., 0., 0., 0.], dtype=dtype), results[0])
+        self.assertEqual(list([3, 0, 2, 6]), list(results[1]))
+
+  def testTopKInfinities(self):
+    """Tests that positive and negative infinity sort correctly."""
+    supported_types = set(
+        [dtypes.bfloat16.as_numpy_dtype, np.float16, np.float32, np.float64])
+    for dtype in supported_types.intersection(self.numeric_types):
+      with self.session() as sess:
+        p = array_ops.placeholder(dtype)
+        with self.test_scope():
+          topk = nn_ops.top_k(p, k=6)
+        results = sess.run(topk, {
+            p:
+                np.array([1, 2, float("inf"), -float("inf"), -1, -2],
+                         dtype=dtype)
+        })
+        self.assertAllEqual(
+            np.array([float("inf"), 2.0, 1.0, -1.0, -2.0, -float("inf")],
+                     dtype=dtype), results[0])
+        self.assertEqual(list([2, 1, 0, 4, 5, 3]), list(results[1]))
+
+  @parameterized.named_parameters(
+      ("Int32", np.int32),
+      ("Int64", np.uint64),
+  )
+  def testInTopK(self, dtype):
+    if dtype in self.numeric_types:
+      array_size = 200 * 1000
+      k_options = [0, 1, 2, 10, 20, 100, 1000, 200 * 1000]
+      batch = 16
+      for x in [np.arange(batch * array_size)]:
+        np.random.shuffle(x)
+        x = np.reshape(x, [batch, array_size])
+        y = np.random.randint(0, array_size, size=batch)
+        for k in k_options:
+          indices = x.argsort(axis=1)[::, -1:-k - 1:-1]
+          expected = [y[i] in indices[i] for i in range(batch)]
+
+          def in_topk(predictions, targets, k=k):
+            return nn_ops.in_top_k(predictions, targets, k)
+
+          self._assertOpOutputMatchesExpected(
+              in_topk,
+              [x.astype(np.float32), y.astype(dtype)],
+              expected=[expected])
+
+
+class SortOpsBenchmark(test.Benchmark):
+  """Microbenchmarks for the sort ops."""
+
+  def _benchmarkSort(self, name, dtype, is_stable, use_xla_jit):
+
+    def get_shuffled_arr(sorted_arr, shape):
+      shuffled = sorted_arr.copy()
+      np.random.shuffle(shuffled)
+      return shuffled.reshape(shape)
+
+    @function.Defun(dtype, dtype)
+    def compare_lt(x1, x2):
+      return x1 < x2
+
+    def builder_fn():
+      shape = (100001,)
+      sorted_arr = np.arange(np.prod(shape), dtype=dtype)
+      shuffled = get_shuffled_arr(sorted_arr, shape)
+      given_result = xla.variadic_sort(
+          [shuffled], dimension=0, is_stable=is_stable, comparator=compare_lt
+      )
+      stable_str = "stable" if is_stable else "unstable"
+      return "%s_%s.shape%s" % (stable_str, name, shape), [given_result]
+
+    xla_test.Benchmark(self, builder_fn, use_xla_jit=use_xla_jit, device="cpu")
+
+  def benchmarkStableSortF16(self):
+    self._benchmarkSort(
+        "sort_f16", dtype=np.float16, is_stable=True, use_xla_jit=False
+    )
+
+  def benchmarkStableSortF32(self):
+    self._benchmarkSort(
+        "sort_f32", dtype=np.float32, is_stable=True, use_xla_jit=False
+    )
+
+  def benchmarkStableSortF64(self):
+    self._benchmarkSort(
+        "sort_f64", dtype=np.float64, is_stable=True, use_xla_jit=False
+    )
+
+  def benchmarkStableSortF16XLA(self):
+    self._benchmarkSort(
+        "sort_f16", dtype=np.float16, is_stable=True, use_xla_jit=True
+    )
+
+  def benchmarkStableSortF32XLA(self):
+    self._benchmarkSort(
+        "sort_f32", dtype=np.float32, is_stable=True, use_xla_jit=True
+    )
+
+  def benchmarkStableSortF64XLA(self):
+    self._benchmarkSort(
+        "sort_f64", dtype=np.float64, is_stable=True, use_xla_jit=True
+    )
+
+  def benchmarkUnstableSortF16(self):
+    self._benchmarkSort(
+        "sort_f16", dtype=np.float16, is_stable=False, use_xla_jit=False
+    )
+
+  def benchmarkUnstableSortF32(self):
+    self._benchmarkSort(
+        "sort_f32", dtype=np.float32, is_stable=False, use_xla_jit=False
+    )
+
+  def benchmarkUnstableSortF64(self):
+    self._benchmarkSort(
+        "sort_f64", dtype=np.float64, is_stable=False, use_xla_jit=False
+    )
+
+  def benchmarkUnstableSortF16XLA(self):
+    self._benchmarkSort(
+        "sort_f16", dtype=np.float16, is_stable=False, use_xla_jit=True
+    )
+
+  def benchmarkUnstableSortF32XLA(self):
+    self._benchmarkSort(
+        "sort_f32", dtype=np.float32, is_stable=False, use_xla_jit=True
+    )
+
+  def benchmarkUnstableSortF64XLA(self):
+    self._benchmarkSort(
+        "sort_f64", dtype=np.float64, is_stable=False, use_xla_jit=True
+    )
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/spacetobatch_op_test.py

@@ -0,0 +1,311 @@
+# 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 SpaceToBatch and BatchToSpace ops."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.platform import test
+
+
+def space_to_batch_direct(input_array, block_shape, paddings):
+  """Direct Python implementation of space-to-batch conversion.
+
+  This is used for tests only.
+
+  Args:
+    input_array: N-D array
+    block_shape: 1-D array of shape [num_block_dims].
+    paddings: 2-D array of shape [num_block_dims, 2].
+
+  Returns:
+    Converted tensor.
+  """
+  input_array = np.array(input_array)
+  block_shape = np.array(block_shape)
+  num_block_dims = len(block_shape)
+  paddings = np.array(paddings).reshape((len(block_shape), 2))
+
+  padded = np.pad(input_array,
+                  pad_width=([[0, 0]] + list(paddings) + [[0, 0]] *
+                             (input_array.ndim - 1 - num_block_dims)),
+                  mode="constant")
+  reshaped_padded_shape = [input_array.shape[0]]
+  output_shape = [input_array.shape[0] * np.prod(block_shape)]
+  for block_dim, block_shape_value in enumerate(block_shape):
+    reduced_size = padded.shape[block_dim + 1] // block_shape_value
+    reshaped_padded_shape.append(reduced_size)
+    output_shape.append(reduced_size)
+    reshaped_padded_shape.append(block_shape_value)
+  reshaped_padded_shape.extend(input_array.shape[num_block_dims + 1:])
+  output_shape.extend(input_array.shape[num_block_dims + 1:])
+
+  reshaped_padded = padded.reshape(reshaped_padded_shape)
+  permuted_reshaped_padded = np.transpose(reshaped_padded, (
+      list(np.arange(num_block_dims) * 2 + 2) + [0] +
+      list(np.arange(num_block_dims) * 2 + 1) + list(
+          np.arange(input_array.ndim - num_block_dims - 1) + 1 + num_block_dims
+          * 2)))
+  return permuted_reshaped_padded.reshape(output_shape)
+
+
+class SpaceToBatchTest(xla_test.XLATestCase):
+  """Tests input-output pairs for the SpaceToBatch and BatchToSpace ops."""
+
+  def _testPad(self, inputs, paddings, block_size, outputs):
+    with self.session() as sess, self.test_scope():
+      for dtype in self.float_types:
+        # outputs = space_to_batch(inputs)
+        placeholder = array_ops.placeholder(dtype)
+        x_tf = gen_array_ops.space_to_batch(
+            placeholder, paddings, block_size=block_size)
+        self.assertAllEqual(sess.run(x_tf, {placeholder: inputs}), outputs)
+        # inputs = batch_to_space(outputs)
+        x_tf = gen_array_ops.batch_to_space(
+            placeholder, paddings, block_size=block_size)
+        self.assertAllEqual(sess.run(x_tf, {placeholder: outputs}), inputs)
+
+  def _testOne(self, inputs, block_size, outputs):
+    paddings = np.zeros((2, 2), dtype=np.int32)
+    self._testPad(inputs, paddings, block_size, outputs)
+
+  # [1, 2, 2, 1] <-> [4, 1, 1, 1]
+  def testSmallInput2x2(self):
+    x_np = [[[[1], [2]], [[3], [4]]]]
+    block_size = 2
+    x_out = [[[[1]]], [[[2]]], [[[3]]], [[[4]]]]
+    self._testOne(x_np, block_size, x_out)
+
+  # [1, 2, 2, 1] <-> [1, 3, 3, 1] (padding) <-> [9, 1, 1, 1]
+  def testSmallInput2x2Pad1x0(self):
+    x_np = [[[[1], [2]], [[3], [4]]]]
+    paddings = np.array([[1, 0], [1, 0]], dtype=np.int32)
+    block_size = 3
+    x_out = [[[[0]]], [[[0]]], [[[0]]], [[[0]]], [[[1]]], [[[2]]], [[[0]]],
+             [[[3]]], [[[4]]]]
+    self._testPad(x_np, paddings, block_size, x_out)
+
+  # Test with depth larger than 1.
+  # [1, 2, 2, 3] <-> [4, 1, 1, 3]
+  def testDepthInput2x2(self):
+    x_np = [[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]]
+    block_size = 2
+    x_out = [[[[1, 2, 3]]], [[[4, 5, 6]]], [[[7, 8, 9]]], [[[10, 11, 12]]]]
+    self._testOne(x_np, block_size, x_out)
+
+  # Test for larger input dimensions.
+  # [1, 4, 4, 1] <-> [4, 2, 2, 1]
+  def testLargerInput2x2(self):
+    x_np = [[[[1], [2], [3], [4]], [[5], [6], [7], [8]],
+             [[9], [10], [11], [12]], [[13], [14], [15], [16]]]]
+    block_size = 2
+    x_out = [[[[1], [3]], [[9], [11]]], [[[2], [4]], [[10], [12]]],
+             [[[5], [7]], [[13], [15]]], [[[6], [8]], [[14], [16]]]]
+    self._testOne(x_np, block_size, x_out)
+
+  # Test with batch larger than 1.
+  # [2, 2, 4, 1] <-> [8, 1, 2, 1]
+  def testBatchInput2x2(self):
+    x_np = [[[[1], [2], [3], [4]], [[5], [6], [7], [8]]],
+            [[[9], [10], [11], [12]], [[13], [14], [15], [16]]]]
+    block_size = 2
+    x_out = [[[[1], [3]]], [[[9], [11]]], [[[2], [4]]], [[[10], [12]]],
+             [[[5], [7]]], [[[13], [15]]], [[[6], [8]]], [[[14], [16]]]]
+    self._testOne(x_np, block_size, x_out)
+
+  # Tests for larger input spatial dimensions AND batch larger than 1, to ensure
+  # that elements are correctly laid out spatially and properly interleaved
+  # along the batch dimension.
+  # [2, 4, 4, 1] <-> [8, 2, 2, 1]
+  def testLargerInputBatch2x2(self):
+    x_np = [[[[1], [2], [3], [4]], [[5], [6], [7], [8]],
+             [[9], [10], [11], [12]], [[13], [14], [15], [16]]],
+            [[[17], [18], [19], [20]], [[21], [22], [23], [24]],
+             [[25], [26], [27], [28]], [[29], [30], [31], [32]]]]
+    x_out = [[[[1], [3]], [[9], [11]]], [[[17], [19]], [[25], [27]]],
+             [[[2], [4]], [[10], [12]]], [[[18], [20]], [[26], [28]]],
+             [[[5], [7]], [[13], [15]]], [[[21], [23]], [[29], [31]]],
+             [[[6], [8]], [[14], [16]]], [[[22], [24]], [[30], [32]]]]
+    block_size = 2
+    self._testOne(x_np, block_size, x_out)
+
+
+class SpaceToBatchNDErrorHandlingTest(xla_test.XLATestCase):
+
+  def testInvalidBlockShape(self):
+    with self.assertRaisesRegex(ValueError, "block_shape must be positive"):
+      with self.session() as sess, self.test_scope():
+        tf_in = constant_op.constant(
+            -3.5e+35, shape=[10, 20, 20], dtype=dtypes.float32)
+        block_shape = constant_op.constant(-10, shape=[2], dtype=dtypes.int64)
+        paddings = constant_op.constant(0, shape=[2, 2], dtype=dtypes.int32)
+        sess.run(array_ops.space_to_batch_nd(tf_in, block_shape, paddings))
+
+  def testOutputSizeOutOfBounds(self):
+    with self.assertRaisesRegex(ValueError,
+                                "Negative.* dimension size caused by overflow"):
+      with self.session() as sess, self.test_scope():
+        tf_in = constant_op.constant(
+            -3.5e+35, shape=[10, 19, 22], dtype=dtypes.float32)
+        block_shape = constant_op.constant(
+            1879048192, shape=[2], dtype=dtypes.int64)
+        paddings = constant_op.constant(0, shape=[2, 2], dtype=dtypes.int32)
+        sess.run(array_ops.space_to_batch_nd(tf_in, block_shape, paddings))
+
+
+class SpaceToBatchNDTest(xla_test.XLATestCase):
+  """Tests input-output pairs for the SpaceToBatchND and BatchToSpaceND ops."""
+
+  def _testPad(self, inputs, block_shape, paddings, outputs):
+    block_shape = np.array(block_shape)
+    paddings = np.array(paddings).reshape((len(block_shape), 2))
+    with self.session() as sess, self.test_scope():
+      for dtype in self.float_types:
+        # TODO(b/68813416): Skip bfloat16's as the input type for direct is
+        # float32 and results in a mismatch, while making testDirect provide the
+        # correctly typed input results in 'no fill-function for data-type'
+        # error.
+        if dtype == dtypes.bfloat16.as_numpy_dtype:
+          continue
+        if dtype == np.float16:
+          actual_inputs = np.array(inputs).astype(dtype)
+          actual_paddings = np.array(paddings).astype(dtype)
+          expected_outputs = np.array(outputs).astype(dtype)
+        else:
+          actual_inputs = inputs
+          actual_paddings = paddings
+          expected_outputs = outputs
+        placeholder = array_ops.placeholder(dtype)
+        # outputs = space_to_batch(inputs)
+        x_tf = array_ops.space_to_batch_nd(placeholder, block_shape,
+                                           actual_paddings)
+        self.assertAllEqual(
+            sess.run(x_tf, {placeholder: actual_inputs}), expected_outputs)
+        # inputs = batch_to_space(outputs)
+        placeholder = array_ops.placeholder(dtype)
+        x_tf = array_ops.batch_to_space_nd(placeholder, block_shape,
+                                           actual_paddings)
+        self.assertAllEqual(
+            sess.run(x_tf, {placeholder: expected_outputs}), actual_inputs)
+
+  def _testDirect(self, input_shape, block_shape, paddings):
+    inputs = np.arange(np.prod(input_shape), dtype=np.float32)
+    inputs = inputs.reshape(input_shape)
+    self._testPad(inputs, block_shape, paddings,
+                  space_to_batch_direct(inputs, block_shape, paddings))
+
+  def testZeroBlockDimsZeroRemainingDims(self):
+    self._testPad(
+        inputs=[1, 2],
+        block_shape=[],
+        paddings=[],
+        outputs=[1, 2],)
+
+  def testZeroBlockDimsOneRemainingDim(self):
+    self._testPad(
+        inputs=[[1, 2], [3, 4]],
+        block_shape=[],
+        paddings=[],
+        outputs=[[1, 2], [3, 4]])
+
+    # Same thing, but with a no-op block dim.
+    self._testPad(
+        inputs=[[1, 2], [3, 4]],
+        block_shape=[1],
+        paddings=[[0, 0]],
+        outputs=[[1, 2], [3, 4]])
+
+  def testZeroBlockDimsTwoRemainingDims(self):
+    self._testPad(
+        inputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+        block_shape=[],
+        paddings=[],
+        outputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+
+    # Same thing, but with a no-op block dim.
+    self._testPad(
+        inputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+        block_shape=[1],
+        paddings=[[0, 0]],
+        outputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+
+    # Same thing, but with two no-op block dims.
+    self._testPad(
+        inputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+        block_shape=[1, 1],
+        paddings=[[0, 0], [0, 0]],
+        outputs=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+
+  def testOneBlockDimZeroRemainingDims(self):
+    self._testPad(
+        inputs=[[1, 2, 3], [4, 5, 6]],
+        block_shape=[2],
+        paddings=[1, 0],
+        outputs=[[0, 2], [0, 5], [1, 3], [4, 6]])
+
+  def testOneBlockDimOneRemainingDim(self):
+    self._testPad(
+        inputs=[[[1, 11], [2, 21], [3, 31]], [[4, 41], [5, 51], [6, 61]]],
+        block_shape=[2],
+        paddings=[1, 0],
+        outputs=[[[0, 0], [2, 21]], [[0, 0], [5, 51]], [[1, 11], [3, 31]],
+                 [[4, 41], [6, 61]]])
+
+  def testDirect0(self):
+    # Test with zero-size remaining dimension.
+    self._testDirect(
+        input_shape=[3, 1, 2, 0], block_shape=[3], paddings=[[0, 2]])
+
+  def testDirect1(self):
+    # Test with zero-size blocked dimension.
+    self._testDirect(
+        input_shape=[3, 0, 2, 5], block_shape=[3], paddings=[[0, 0]])
+
+  def testDirect2(self):
+    # Test with padding up from zero size.
+    self._testDirect(
+        input_shape=[3, 0, 2, 5], block_shape=[3], paddings=[[1, 2]])
+
+  def testDirect3(self):
+    self._testDirect(
+        input_shape=[3, 3, 4, 5, 2],
+        block_shape=[3, 4, 2],
+        paddings=[[1, 2], [0, 0], [3, 0]])
+
+  def testDirect4(self):
+    self._testDirect(
+        input_shape=[3, 3, 4, 5, 2],
+        block_shape=[3, 4, 2, 2],
+        paddings=[[1, 2], [0, 0], [3, 0], [0, 0]])
+
+  def testDirect5(self):
+    self._testDirect(
+        input_shape=[3, 2, 2, 3, 4, 5, 2, 5],
+        block_shape=[1, 1, 3, 4, 2, 2],
+        paddings=[[0, 0], [0, 0], [1, 2], [0, 0], [3, 0], [0, 0]])
+
+  def testDirect6(self):
+    self._testDirect(
+        input_shape=[3, 2, 2, 3, 4, 5, 2, 5],
+        block_shape=[1, 1, 3, 4, 2, 2, 1],
+        paddings=[[0, 0], [0, 0], [1, 2], [0, 0], [3, 0], [0, 0], [0, 0]])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/sparse_to_dense_op_test.py

@@ -0,0 +1,124 @@
+# Copyright 2015 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 tensorflow.kernels.sparse_op."""
+
+import numpy as np
+
+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 sparse_ops
+from tensorflow.python.platform import test
+
+
+def _SparseToDense(sparse_indices,
+                   output_size,
+                   sparse_values,
+                   default_value,
+                   validate_indices=True):
+  feed_sparse_indices = array_ops.placeholder(dtypes.int32)
+  feed_dict = {feed_sparse_indices: sparse_indices}
+  return sparse_ops.sparse_to_dense(
+      feed_sparse_indices,
+      output_size,
+      sparse_values,
+      default_value=default_value,
+      validate_indices=validate_indices).eval(feed_dict=feed_dict)
+
+
+class SparseToDenseTest(xla_test.XLATestCase):
+
+  def testInt(self):
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([1, 3], [5], 1, 0)
+    np_ans = np.array([0, 1, 0, 1, 0]).astype(np.int32)
+    self.assertAllClose(np_ans, tf_ans)
+
+  def testFloat(self):
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([1, 3], [5], 1.0, 0.0)
+    np_ans = np.array([0, 1, 0, 1, 0]).astype(np.float32)
+    self.assertAllClose(np_ans, tf_ans)
+
+  def testSetValue(self):
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([1, 3], [5], [1, 2], -1)
+    np_ans = np.array([-1, 1, -1, 2, -1]).astype(np.int32)
+    self.assertAllClose(np_ans, tf_ans)
+
+  def testSetSingleValue(self):
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([1, 3], [5], 1, -1)
+    np_ans = np.array([-1, 1, -1, 1, -1]).astype(np.int32)
+    self.assertAllClose(np_ans, tf_ans)
+
+  def test2d(self):
+    # pylint: disable=bad-whitespace
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([[1, 3], [2, 0]], [3, 4], 1, -1)
+    np_ans = np.array([[-1, -1, -1, -1],
+                       [-1, -1, -1,  1],
+                       [ 1, -1, -1, -1]]).astype(np.int32)
+    self.assertAllClose(np_ans, tf_ans)
+
+  def testZeroDefault(self):
+    with self.session():
+      x = sparse_ops.sparse_to_dense(2, [4], 7).eval()
+      self.assertAllEqual(x, [0, 0, 7, 0])
+
+  def test3d(self):
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([[1, 3, 0], [2, 0, 1]], [3, 4, 2], 1, -1)
+    np_ans = np.ones((3, 4, 2), dtype=np.int32) * -1
+    np_ans[1, 3, 0] = 1
+    np_ans[2, 0, 1] = 1
+    self.assertAllClose(np_ans, tf_ans)
+
+  def testDegenerateIndexMatrix(self):
+    with self.session(), self.test_scope():
+      tf_ans = _SparseToDense([[2], [3], [4], [5], [6], [7], [8], [9]], [10],
+                              [1, 2, 3, 4, 5, 6, 7, 8], -1)
+    self.assertAllClose([-1, -1, 1, 2, 3, 4, 5, 6, 7, 8], tf_ans)
+
+  def testBadShape(self):
+    with self.session(), self.test_scope():
+      with self.assertRaisesWithPredicateMatch(ValueError, "must be rank 1"):
+        _SparseToDense([1, 3], [[5], [3]], 1, -1)
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testBadValue(self):
+    with self.session(), self.test_scope():
+      with self.assertRaisesOpError(
+          r"sparse_values has incorrect shape \[2,1\], "
+          r"should be \[\] or \[2\]"):
+        _SparseToDense([1, 3], [5], [[5], [3]], -1)
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testBadNumValues(self):
+    with self.session(), self.test_scope():
+      with self.assertRaisesOpError(
+          r"sparse_values has incorrect shape \[3\], should be \[\] or \[2\]"):
+        _SparseToDense([1, 3], [5], [1, 2, 3], -1)
+
+  @test_util.disable_mlir_bridge("Error handling")
+  def testBadDefault(self):
+    with self.session(), self.test_scope():
+      with self.assertRaisesOpError("default_value should be a scalar"):
+        _SparseToDense([1, 3], [5], [1, 2], [0])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/special_math_test.py

@@ -0,0 +1,624 @@
+# 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 special math operations."""
+
+import os
+
+from absl import flags
+from absl.testing import parameterized
+
+import numpy as np
+import scipy.special as sps
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import gen_random_ops
+from tensorflow.python.ops import gradient_checker_v2
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+flags.DEFINE_bool('vary_seed', False,
+                  ('Whether to vary the PRNG seed unpredictably.  '
+                   'With --runs_per_test=N, produces N iid runs.'))
+
+NUM_SAMPLES = int(1e3)
+
+
+@def_function.function(jit_compile=True)
+def _igamma(a, x):
+  return math_ops.igamma(a, x)
+
+
+@def_function.function(jit_compile=True)
+def _igammac(a, x):
+  return math_ops.igammac(a, x)
+
+
+@def_function.function(jit_compile=True)
+def _polygamma(n, x):
+  return math_ops.polygamma(n, x)
+
+
+@def_function.function(jit_compile=True)
+def _zeta(a, q):
+  return math_ops.zeta(a, q)
+
+
+# This is df/da / df/dx, where f = igamma.
+def implicit_reparameterization_grad(a, x):
+  log_prob = math_ops.xlogy(a - 1., x) - math_ops.lgamma(a) - x
+  prob = math_ops.exp(log_prob)
+  return -gen_math_ops.igamma_grad_a(a, x) / prob
+
+
+@def_function.function(jit_compile=True)
+def _log1p(x):
+  return math_ops.log1p(x)
+
+
+class Log1pTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def setUp(self):
+    if flags.FLAGS.vary_seed:
+      entropy = os.urandom(64)
+      answer = int.from_bytes(entropy, 'big')
+      np.random.seed(answer % (2**32 - 1))
+    super(Log1pTest, self).setUp()
+
+  def adjust_tolerance(self, dtype, rtol, atol):
+    if self.device in ['TPU']:
+      if dtype == np.float32:
+        return 4e-4, 0.0
+      return 1e-10, 0.0
+    if self.device in ['XLA_GPU', 'GPU']:
+      if dtype == np.float32:
+        return max(rtol, 2.5e-07), atol
+    return rtol, atol
+
+  def _test_range(self, low, high, dtype, rtol, atol, is_negative=False):
+    # Test values near zero.
+    rtol, atol = self.adjust_tolerance(dtype, rtol, atol)
+    x = np.exp(np.random.uniform(
+        low=low, high=high, size=[NUM_SAMPLES])).astype(dtype)
+    if is_negative:
+      x = -x
+    expected_values = np.log1p(x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = _log1p(x)
+      actual = sess.run(actual)
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testSmallX(self, dtype, rtol, atol):
+    self._test_range(-40., -20., dtype, rtol, atol, is_negative=False)
+    self._test_range(-40., -20., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 2e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testGreaterThanNegativeTwentyExponent(self, dtype, rtol, atol):
+    self._test_range(-20., -10., dtype, rtol, atol, is_negative=False)
+    self._test_range(-20., -10., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 2e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testGreaterThanNegativeTenExponent(self, dtype, rtol, atol):
+    self._test_range(-10., -5., dtype, rtol, atol, is_negative=False)
+    self._test_range(-10., -5., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 2e-7, 0.),
+                            (np.float64, 1e-15, 0.))
+  def testGreaterThanNegativeFiveExponent(self, dtype, rtol, atol):
+    self._test_range(-5., -1., dtype, rtol, atol, is_negative=False)
+    self._test_range(-5., -1., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 4e-7, 0.),
+                            (np.float64, 3e-14, 0.))
+  def testXGreaterThanOneTenth(self, dtype, rtol, atol):
+    self._test_range(-1., 0., dtype, rtol, atol, is_negative=False)
+    self._test_range(-1., 0., dtype, rtol, atol, is_negative=True)
+
+  @parameterized.parameters((np.float32, 2e-7, 0.),
+                            (np.float64, 2e-15, 0.))
+  def testXGreaterThanOne(self, dtype, rtol, atol):
+    self._test_range(0., 3., dtype, rtol, atol, is_negative=False)
+
+
+class ZetaTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def setUp(self):
+    if flags.FLAGS.vary_seed:
+      entropy = os.urandom(64)
+      answer = int.from_bytes(entropy, 'big')
+      np.random.seed(answer % (2**32 - 1))
+    super(ZetaTest, self).setUp()
+
+  def adjust_tolerance_for_tpu(self, dtype, rtol, atol):
+    if self.device not in ['TPU']:
+      return rtol, atol
+
+    if dtype == np.float32:
+      return 2e-2, 1e-7
+    return 2e-4, 1e-20
+
+  def testBadValues(self):
+    q = np.random.uniform(low=0.3, high=20., size=[10])
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta(np.float64(1.), q)
+      actual = sess.run(y)
+    # When x == 1, this is the Harmonic series.
+    self.assertTrue(np.all(np.isinf(actual)))
+
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta(np.float64(0.1), q)
+      actual = sess.run(y)
+    # When x < 1, this is undefined.
+    self.assertTrue(np.all(np.isnan(actual)))
+
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta([1.1, 1.2, 2.1, 2.2, 3.1], [-2.0, -1.1, -1.0, -0.5, -0.1])
+      actual = sess.run(y)
+    # For q <= 0, x must be an integer.
+    self.assertTrue(np.all(np.isnan(actual)))
+
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta([2.0, 4.0, 6.0], [0.0, -1.0, -2.0])
+      actual = sess.run(y)
+    # For integer q <= 0, zeta has poles with a defined limit of +inf where x is
+    # an even integer.
+    self.assertTrue(np.all(np.isinf(actual)))
+
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta([3.0, 5.0, 7.0], [0.0, -1.0, -2.0])
+      actual = sess.run(y)
+    # For non-positive integer q, zeta has poles with an undefined limit where x
+    # is an odd integer.
+    self.assertTrue(np.all(np.isnan(actual)))
+
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta([1.1, 2.2, 3.3], [-1.1, -1.0, 0.0])
+      actual = sess.run(y)
+    # For non-positive q, zeta is not defined if x is not an integer.
+    self.assertTrue(np.all(np.isnan(actual)))
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testLargeXSmallQ(self, dtype, rtol, atol):
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    if self.device not in ['XLA_GPU', 'XLA_CPU'] and dtype == np.float64:
+      # TODO(b/165739664): Figure out why on TPU F64 Zeta sometimes returns
+      # infs.
+      self.skipTest(
+          'Skipping test because some F64 operations are numerically '
+          'unstable on TPU.')
+
+    x = np.random.uniform(low=100., high=200., size=[NUM_SAMPLES]).astype(dtype)
+    q = np.random.uniform(low=0.3, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.zeta(x, q)
+    with self.session() as sess:
+      with self.test_scope():
+        y = _zeta(x, q)
+      actual = sess.run(y)
+
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testSmallValues(self, dtype, rtol, atol):
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(low=1.1, high=10., size=[NUM_SAMPLES]).astype(dtype)
+    q = np.random.uniform(
+        low=np.finfo(dtype).tiny, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.zeta(x, q)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_zeta(x, q))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testMediumValues(self, dtype, rtol, atol):
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    x = np.random.uniform(low=1.1, high=100., size=[NUM_SAMPLES]).astype(dtype)
+    q = np.random.uniform(low=1., high=1e1, size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.zeta(x, q)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_zeta(x, q))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 2e-2, 1e-5), (np.float64, 1e-4, 1e-30))
+  def testLargeValues(self, dtype, rtol, atol):
+    x = np.random.uniform(
+        low=100., high=int(1e3), size=[NUM_SAMPLES]).astype(dtype)
+    q = np.random.uniform(
+        low=1., high=int(1e1), size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.zeta(x, q)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_zeta(x, q))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+
+class PolygammaTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def setUp(self):
+    if flags.FLAGS.vary_seed:
+      entropy = os.urandom(64)
+      answer = int.from_bytes(entropy, 'big')
+      np.random.seed(answer % (2**32 - 1))
+    super(PolygammaTest, self).setUp()
+
+  def adjust_tolerance_for_tpu(self, dtype, rtol, atol):
+    if self.device not in ['TPU']:
+      return rtol, atol
+
+    if dtype == np.float32:
+      return 2e-2, 1e-7
+    return 2e-4, 1e-20
+
+  def testBadValues(self):
+    x = np.random.uniform(low=0.3, high=20., size=[10])
+    with self.session() as sess:
+      with self.test_scope():
+        y = _polygamma(np.float64(-1.), x)
+      actual = sess.run(y)
+    # Not defined for negative numbers.
+    self.assertTrue(np.all(np.isnan(actual)))
+
+    with self.session() as sess:
+      with self.test_scope():
+        y = _polygamma(np.float64(0.1), x)
+      actual = sess.run(y)
+    # Not defined for non-integers.
+    self.assertTrue(np.all(np.isnan(actual)))
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testRecoverDigamma(self, dtype, rtol, atol):
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    if self.device not in ['XLA_GPU', 'XLA_CPU'] and dtype == np.float64:
+      self.skipTest(
+          'Skipping test because some F64 operations are '
+          'numerically unstable on TPU.'
+      )
+
+    x = np.random.uniform(low=0.1, high=50., size=[NUM_SAMPLES]).astype(dtype)
+    expected_values = sps.digamma(x)
+    with self.session() as sess:
+      with self.test_scope():
+        y = _polygamma(dtype(0.), x)
+      actual = sess.run(y)
+
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testSmallN(self, dtype, rtol, atol):
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    n = np.random.randint(low=1, high=5, size=[NUM_SAMPLES]).astype(dtype)
+    x = np.random.uniform(
+        low=np.finfo(dtype).tiny, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.polygamma(n, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_polygamma(n, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testMediumLargeN(self, dtype, rtol, atol):
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    n = np.random.randint(low=5, high=10, size=[NUM_SAMPLES]).astype(dtype)
+    x = np.random.uniform(low=1., high=1e1, size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.polygamma(n, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_polygamma(n, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+
+class IgammaTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def setUp(self):
+    if flags.FLAGS.vary_seed:
+      entropy = os.urandom(64)
+      answer = int.from_bytes(entropy, 'big')
+      np.random.seed(answer % (2**32 - 1))
+    super(IgammaTest, self).setUp()
+
+  # Skip Float64 test on TPU due to missing ops.
+  def maybe_skip_test(self, dtype):
+    if self.device not in ['XLA_GPU', 'XLA_CPU'] and dtype == np.float64:
+      self.skipTest(
+          'Skipping test because some F64 operations not supported on TPU.')
+
+  def adjust_tolerance_for_tpu(self, dtype, rtol, atol):
+    if self.device not in ['TPU']:
+      return rtol, atol
+
+    if dtype == np.float32:
+      return 2e-2, 1e-7
+    return 2e-4, 1e-20
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testLargeXSmallA(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(low=100., high=200., size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(low=0.3, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammainc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        y = _igamma(a, x)
+      actual = sess.run(y)
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testSmallValues(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(
+        low=np.finfo(dtype).tiny, high=1., size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(
+        low=np.finfo(dtype).tiny, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammainc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_igamma(a, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testMediumValues(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(low=1., high=100., size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(low=1., high=100., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammainc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_igamma(a, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 2e-2, 1e-5), (np.float64, 1e-4, 1e-30))
+  def testLargeValues(self, dtype, rtol, atol):
+    if self.device == 'TPU':
+      # TODO(b/154908275): Remove this once fixed for large a, x.
+      self.skipTest('Skipping test since numerically unstable on TPU.')
+    # Test values near zero.
+    x = np.random.uniform(
+        low=100., high=int(1e4), size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(
+        low=100., high=int(1e4), size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammainc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_igamma(a, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  # We don't check small values because the numerical gradients become quite
+  # large.
+  @parameterized.parameters((np.float32, 0.09), (np.float64, 1e-7))
+  def testGradMediumValues(self, dtype, tolerance):
+    self.maybe_skip_test(dtype)
+    with self.session():
+      with self.test_scope():
+        x = constant_op.constant(
+            np.random.uniform(low=1., high=100.,
+                              size=[NUM_SAMPLES]).astype(dtype))
+        a = constant_op.constant(
+            np.random.uniform(low=1., high=100.,
+                              size=[NUM_SAMPLES]).astype(dtype))
+
+        f = lambda b: _igamma(b, x)
+        max_error = gradient_checker_v2.max_error(
+            *gradient_checker_v2.compute_gradient(f, x=[a], delta=1e-3))
+    self.assertLessEqual(max_error, tolerance)
+
+  @parameterized.parameters((np.float32, 0.5), (np.float64, 1e-7))
+  def testGradLargeValues(self, dtype, tolerance):
+    self.maybe_skip_test(dtype)
+    with self.session():
+      with self.test_scope():
+        x = constant_op.constant(
+            np.random.uniform(low=100., high=int(1e4),
+                              size=[NUM_SAMPLES]).astype(dtype))
+        a = constant_op.constant(
+            np.random.uniform(low=100., high=int(1e4),
+                              size=[NUM_SAMPLES]).astype(dtype))
+
+        f = lambda b: _igamma(b, x)
+        max_error = gradient_checker_v2.max_error(
+            *gradient_checker_v2.compute_gradient(f, x=[a], delta=1e-2))
+    self.assertLessEqual(max_error, tolerance)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testRandomGammaGradSmallValues(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+
+    with self.session() as sess:
+      with self.test_scope():
+        x = constant_op.constant(
+            np.random.uniform(
+                low=np.finfo(dtype).tiny, high=1.,
+                size=[NUM_SAMPLES]).astype(dtype))
+        a = constant_op.constant(
+            np.random.uniform(
+                low=np.finfo(dtype).tiny, high=1.,
+                size=[NUM_SAMPLES]).astype(dtype))
+        gamma_sample_grad = gen_random_ops.random_gamma_grad(a, x)
+        actual_grad = implicit_reparameterization_grad(a, x)
+        gamma_sample_grad, actual_grad = sess.run(
+            [gamma_sample_grad, actual_grad])
+        # We do this because the ratio computed in
+        # implicit_reparameterization_grad can very easily result in a NaN due
+        # to the computed numerator and denominator zeroing out.
+        gamma_sample_grad = gamma_sample_grad[
+            ~np.logical_or(np.isnan(actual_grad), np.isinf(actual_grad))]
+        actual_grad = actual_grad[
+            ~np.logical_or(np.isnan(actual_grad), np.isinf(actual_grad))]
+    self.assertAllClose(actual_grad, gamma_sample_grad, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testRandomGammaGradMediumValues(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+
+    with self.session() as sess:
+      with self.test_scope():
+        x = constant_op.constant(
+            np.random.uniform(low=1., high=10.,
+                              size=[NUM_SAMPLES]).astype(dtype))
+        a = constant_op.constant(
+            np.random.uniform(low=1., high=10.,
+                              size=[NUM_SAMPLES]).astype(dtype))
+        gamma_sample_grad = gen_random_ops.random_gamma_grad(a, x)
+        actual_grad = implicit_reparameterization_grad(a, x)
+        gamma_sample_grad, actual_grad = sess.run(
+            [gamma_sample_grad, actual_grad])
+        # We do this because the ratio computed in
+        # implicit_reparameterization_grad can very easily result in a NaN due
+        # to the computed numerator and denominator zeroing out.
+        gamma_sample_grad = gamma_sample_grad[
+            ~np.logical_or(np.isnan(actual_grad), np.isinf(actual_grad))]
+        actual_grad = actual_grad[
+            ~np.logical_or(np.isnan(actual_grad), np.isinf(actual_grad))]
+    self.assertAllClose(actual_grad, gamma_sample_grad, atol=atol, rtol=rtol)
+
+
+class IgammacTest(xla_test.XLATestCase, parameterized.TestCase):
+
+  def setUp(self):
+    if flags.FLAGS.vary_seed:
+      entropy = os.urandom(64)
+      answer = int.from_bytes(entropy, 'big')
+      np.random.seed(answer % (2**32 - 1))
+    super(IgammacTest, self).setUp()
+
+  # Skip Float64 test on TPU due to missing ops.
+  def maybe_skip_test(self, dtype):
+    if self.device not in ['XLA_GPU', 'XLA_CPU'] and dtype == np.float64:
+      # TODO(b/154908275): Remove this once fixed for large a, x.
+      self.skipTest(
+          'Skipping test because some F64 operations not supported on TPU.')
+
+  def adjust_tolerance_for_tpu(self, dtype, rtol, atol):
+    if self.device not in ['TPU']:
+      return rtol, atol
+
+    if dtype == np.float32:
+      return 2e-2, 1e-7
+    return 2e-4, 1e-20
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testLargeXSmallA(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(low=100., high=200., size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(low=0.3, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammaincc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        y = _igammac(a, x)
+      actual = sess.run(y)
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testSmallValues(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(
+        low=np.finfo(dtype).tiny, high=1., size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(
+        low=np.finfo(dtype).tiny, high=1., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammaincc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_igammac(a, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 1e-2, 1e-11),
+                            (np.float64, 1e-4, 1e-30))
+  def testMediumValues(self, dtype, rtol, atol):
+    self.maybe_skip_test(dtype)
+    rtol, atol = self.adjust_tolerance_for_tpu(dtype, rtol, atol)
+    # Test values near zero.
+    x = np.random.uniform(low=1., high=100., size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(low=1., high=100., size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammaincc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_igammac(a, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+  @parameterized.parameters((np.float32, 2e-2, 1e-5), (np.float64, 1e-4, 1e-30))
+  def testLargeValues(self, dtype, rtol, atol):
+    if self.device == 'TPU':
+      self.skipTest('Skipping test since numerically unstable on TPU.')
+    # Test values near zero.
+    x = np.random.uniform(
+        low=100., high=int(1e4), size=[NUM_SAMPLES]).astype(dtype)
+    a = np.random.uniform(
+        low=100., high=int(1e4), size=[NUM_SAMPLES]).astype(dtype)
+
+    expected_values = sps.gammaincc(a, x)
+    with self.session() as sess:
+      with self.test_scope():
+        actual = sess.run(_igammac(a, x))
+    self.assertAllClose(expected_values, actual, atol=atol, rtol=rtol)
+
+
+if __name__ == '__main__':
+  os.environ['XLA_FLAGS'] = '--xla_cpu_enable_fast_math=false'
+  test.main()

tensorflow/compiler/tests/stack_ops_test.py

@@ -0,0 +1,121 @@
+# 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.
+# ==============================================================================
+"""Tests for tensorflow.ops.stack_ops."""
+
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.compiler.xla import xla
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_data_flow_ops
+from tensorflow.python.platform import test
+
+
+class StackOpTest(xla_test.XLATestCase):
+
+  def testStackPushPop(self):
+    with self.session(), self.test_scope():
+
+      v = array_ops.placeholder(dtypes.float32)
+
+      def fn():
+        h = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+        c = gen_data_flow_ops.stack_push_v2(h, v)
+        with ops.control_dependencies([c]):
+          c1 = gen_data_flow_ops.stack_pop_v2(h, dtypes.float32)
+        return c1
+
+      self.assertAllClose([[4.0, 5.0]],
+                          xla.compile(fn)[0].eval({v: [[4.0, 5.0]]}))
+
+  def testStackPushPopSwap(self):
+    with self.session(), self.test_scope():
+      a = np.arange(2000)
+      x = array_ops.placeholder(dtypes.float32)
+
+      def fn():
+        h = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+        c = gen_data_flow_ops.stack_push_v2(h, x, swap_memory=True)
+        with ops.control_dependencies([c]):
+          return gen_data_flow_ops.stack_pop_v2(h, dtypes.float32)
+
+      self.assertAllClose(a, xla.compile(fn)[0].eval({x: a}))
+
+  def testMultiStack(self):
+    with self.session(), self.test_scope():
+      v = array_ops.placeholder(dtypes.float32)
+
+      def fn():
+        h1 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+        c1 = gen_data_flow_ops.stack_push_v2(h1, v)
+        with ops.control_dependencies([c1]):
+          c1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32)
+        h2 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="bar")
+        c2 = gen_data_flow_ops.stack_push_v2(h2, 5.0)
+        with ops.control_dependencies([c2]):
+          c2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32)
+        return c1 + c2
+
+      self.assertAllClose(9.0, xla.compile(fn)[0].eval({v: 4.0}))
+
+  def testSameNameStacks(self):
+    """Different stacks with the same name do not interfere."""
+    with self.session() as sess, self.test_scope():
+      v1 = array_ops.placeholder(dtypes.float32)
+      v2 = array_ops.placeholder(dtypes.float32)
+
+      def fn():
+        h1 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+        h2 = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+
+        c1 = gen_data_flow_ops.stack_push_v2(h1, v1)
+        with ops.control_dependencies([c1]):
+          c2 = gen_data_flow_ops.stack_push_v2(h2, v2)
+        with ops.control_dependencies([c2]):
+          pop1 = gen_data_flow_ops.stack_pop_v2(h1, dtypes.float32)
+          pop2 = gen_data_flow_ops.stack_pop_v2(h2, dtypes.float32)
+        return [pop1, pop2]
+
+      [pop1_compiled, pop2_compiled] = xla.compile(fn)
+      out1, out2 = sess.run([pop1_compiled, pop2_compiled], {v1: 4.0, v2: 5.0})
+      self.assertAllClose(out1, 4.0)
+      self.assertAllClose(out2, 5.0)
+
+  def testCloseStack(self):
+    with self.session() as sess, self.test_scope():
+
+      def fn():
+        h = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+        gen_data_flow_ops.stack_close_v2(h)
+
+      sess.run(xla.compile(fn))
+
+  def testPushCloseStack(self):
+    with self.session() as sess, self.test_scope():
+      v = array_ops.placeholder(dtypes.float32)
+
+      def fn():
+        h = gen_data_flow_ops.stack_v2(5, dtypes.float32, stack_name="foo")
+        c = gen_data_flow_ops.stack_push_v2(h, v)
+        with ops.control_dependencies([c]):
+          gen_data_flow_ops.stack_close_v2(h)
+
+      sess.run(xla.compile(fn), {v: [[4.0, 5.0]]})
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/compiler/tests/stateful_random_ops_test.py

@@ -0,0 +1,389 @@
+# 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 stateful random-number generation ops."""
+
+import itertools
+import os
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.client import device_lib
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import config
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.kernel_tests.random import util as \
+random_test_util
+from tensorflow.python.ops import gen_stateful_random_ops
+from tensorflow.python.ops import random_ops_util
+from tensorflow.python.ops import stateful_random_ops as \
+random
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import flags
+from tensorflow.python.platform import test
+
+
+FLAGS = flags.FLAGS
+
+
+def xla_device():
+  devices = device_lib.list_local_devices()
+  def find_type(device_type):
+    for d in devices:
+      if d.device_type == device_type:
+        return d
+    return None
+  d = find_type("TPU") or find_type("XLA_GPU") or find_type("XLA_CPU")
+  if d is None:
+    raise ValueError(
+        "Can't find any XLA device. Available devices:\n%s" % devices)
+  return d
+
+
+def xla_device_name():
+  return str(xla_device().name)
+
+
+ALGS = [
+    random_ops_util.Algorithm.PHILOX.value,
+    random_ops_util.Algorithm.THREEFRY.value,
+    random_ops_util.Algorithm.AUTO_SELECT.value,
+]
+INTS = [dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64]
+FLOATS = [dtypes.bfloat16, dtypes.float32, dtypes.float64]
+
+
+class StatefulRandomOpsTest(xla_test.XLATestCase, parameterized.TestCase):
+  """Test cases for stateful random-number generator operators."""
+
+  @parameterized.parameters(ALGS)
+  def testSimple(self, alg):
+    """A simple test."""
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=0, alg=alg)
+      gen.normal(shape=(3,))
+      gen.uniform(shape=(3,), minval=0, maxval=10, dtype=dtypes.uint32)
+      gen.uniform_full_int(shape=(3,))
+
+  @parameterized.parameters(ALGS)
+  def testDefun(self, alg):
+    """Test for defun."""
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=0, alg=alg)
+      @def_function.function
+      def f():
+        x = gen.normal(shape=(3,))
+        y = gen.uniform(shape=(3,), minval=0, maxval=10, dtype=dtypes.uint32)
+        z = gen.uniform_full_int(shape=(3,))
+        return (x, y, z)
+      f()
+
+  def _compareToKnownOutputs(self, g, counter, key, expect):
+    """Compares against known outputs for specific counter and key inputs."""
+    def uint32s_to_uint64(a, b):
+      return b << 32 | a
+
+    def uint32s_to_uint64s(ls):
+      return [uint32s_to_uint64(ls[2 * i], ls[2 * i + 1])
+              for i in range(len(ls) // 2)]
+
+    ctr_len = len(counter)
+    counter = uint32s_to_uint64s(counter)
+    key = uint32s_to_uint64s(key)
+    state = counter + key
+    g.reset(state)
+    got = g.uniform_full_int(shape=(ctr_len,), dtype=dtypes.uint32)
+    self.assertAllEqual(expect, got)
+    g.reset(state)
+    got = g.uniform_full_int(shape=(ctr_len // 2,), dtype=dtypes.uint64)
+    self.assertAllEqual(uint32s_to_uint64s(expect), got)
+
+  def testThreefry2x32(self):
+    """Tests ThreeFry2x32 conforms to known results.
+    """
+    # Based on
+    # https://github.com/google/jax/blob/8565a3486adf16beb388b2364c9cd930d7a0d92d/tests/random_test.py#L65-L85
+    # which is in turn based on
+    # https://github.com/DEShawResearch/Random123-Boost/blob/65e3d874b67aa7b3e02d5ad8306462f52d2079c0/libs/random/test/test_threefry.cpp#L30-L32
+
+    with ops.device(xla_device_name()):
+      g = random.Generator.from_seed(seed=0, alg=random.RNG_ALG_THREEFRY)
+      self._compareToKnownOutputs(
+          g,
+          [0x00000000, 0x00000000], [0x00000000, 0x00000000],
+          [0x6b200159, 0x99ba4efe])
+      self._compareToKnownOutputs(
+          g,
+          [0xffffffff, 0xffffffff], [0xffffffff, 0xffffffff],
+          [0x1cb996fc, 0xbb002be7])
+      self._compareToKnownOutputs(
+          g,
+          [0x243f6a88, 0x85a308d3], [0x13198a2e, 0x03707344],
+          [0xc4923a9c, 0x483df7a0])
+
+  def testPhilox4x32(self):
+    """Tests Philox4x32 conforms to known results.
+    """
+    # Based on
+    # https://github.com/DEShawResearch/Random123-Boost/blob/65e3d874b67aa7b3e02d5ad8306462f52d2079c0/libs/random/test/test_philox.cpp#L50-L52
+
+    with ops.device(xla_device_name()):
+      g = random.Generator.from_seed(seed=0, alg=random.RNG_ALG_PHILOX)
+      self._compareToKnownOutputs(
+          g,
+          [0x00000000, 0x00000000, 0x00000000, 0x00000000],
+          [0x00000000, 0x00000000],
+          [0x6627e8d5, 0xe169c58d, 0xbc57ac4c, 0x9b00dbd8])
+      self._compareToKnownOutputs(
+          g,
+          [0xffffffff, 0xffffffff, 0xffffffff, 0xffffffff],
+          [0xffffffff, 0xffffffff],
+          [0x408f276d, 0x41c83b0e, 0xa20bc7c6, 0x6d5451fd])
+      self._compareToKnownOutputs(
+          g,
+          [0x243f6a88, 0x85a308d3, 0x13198a2e, 0x03707344],
+          [0xa4093822, 0x299f31d0],
+          [0xd16cfe09, 0x94fdcceb, 0x5001e420, 0x24126ea1])
+
+  @parameterized.parameters(INTS)
+  def testXLAEqualsCPU(self, dtype):
+    """Tests that XLA and CPU kernels generate the same integers."""
+    seed = 1234
+    shape = [315, 49]
+    with ops.device("/device:CPU:0"):
+      cpu_gen = random.Generator.from_seed(seed=seed, alg=random.RNG_ALG_PHILOX)
+    with ops.device(xla_device_name()):
+      xla_gen = random.Generator.from_seed(seed=seed, alg=random.RNG_ALG_PHILOX)
+    # Repeat multiple times to make sure that the state after
+    # number-generation are the same between CPU and XLA.
+    for _ in range(5):
+      with ops.device("/device:CPU:0"):
+        # Test both number-generation and skip
+        cpu = cpu_gen.uniform_full_int(shape=shape, dtype=dtype)
+        cpu_gen.skip(100)
+      with ops.device(xla_device_name()):
+        xla = xla_gen.uniform_full_int(shape=shape, dtype=dtype)
+        xla_gen.skip(100)
+      self.assertAllEqual(cpu, xla)
+      self.assertAllEqual(cpu_gen.state, xla_gen.state)
+
+  def testXLAEqualsCPUAroundCounterOverflow(self):
+    """Tests XLA and CPU kernels generate the same integers in overflow case.
+
+       Specifically this tests the case where the counter is incremented past
+       what can fit within 64 bits of the 128 bit Philox counter.
+    """
+    dtype = dtypes.uint64
+    seed = 2**64 - 10
+    shape = [315, 49]
+    with ops.device("/device:CPU:0"):
+      cpu_gen = random.Generator.from_seed(seed=seed, alg=random.RNG_ALG_PHILOX)
+    with ops.device(xla_device_name()):
+      xla_gen = random.Generator.from_seed(seed=seed, alg=random.RNG_ALG_PHILOX)
+    # Repeat multiple times to make sure that the state after
+    # number-generation are the same between CPU and XLA.
+    for _ in range(5):
+      with ops.device("/device:CPU:0"):
+        # Test both number-generation and skip
+        cpu = cpu_gen.uniform_full_int(shape=shape, dtype=dtype)
+        cpu_gen.skip(100)
+      with ops.device(xla_device_name()):
+        xla = xla_gen.uniform_full_int(shape=shape, dtype=dtype)
+        xla_gen.skip(100)
+      self.assertAllEqual(cpu, xla)
+      self.assertAllEqual(cpu_gen.state, xla_gen.state)
+    self.assertAllEqual(cpu, xla)
+
+  def _testRngIsNotConstant(self, rng, dtype):
+    # Tests that 'rng' does not always return the same value.
+    # The random-number generator, if working correctly, should produce the
+    # same output multiple times with low probability.
+    x = rng(dtype).numpy()
+    y = rng(dtype).numpy()
+    self.assertFalse(np.array_equal(x, y))
+
+  def check_dtype(self, dtype):
+    device = xla_device()
+    if device.device_type == "TPU" and dtype == dtypes.float64:
+      self.skipTest("TPU doesn't support float64.")
+
+  @parameterized.parameters(list(itertools.product(ALGS, INTS + FLOATS)))
+  def testUniformIsNotConstant(self, alg, dtype):
+    self.check_dtype(dtype)
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=1234, alg=alg)
+      def rng(dtype):
+        maxval = dtype.max
+        return gen.uniform(shape=[2], dtype=dtype, maxval=maxval)
+
+      self._testRngIsNotConstant(rng, dtype)
+
+  @parameterized.parameters(list(itertools.product(ALGS, FLOATS)))
+  def testNormalIsNotConstant(self, alg, dtype):
+    self.check_dtype(dtype)
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=1234, alg=alg)
+      def rng(dtype):
+        return gen.normal(shape=[2], dtype=dtype)
+
+      self._testRngIsNotConstant(rng, dtype)
+
+  @parameterized.parameters(list(itertools.product(ALGS, INTS + FLOATS)))
+  def testUniformIsInRange(self, alg, dtype):
+    self.check_dtype(dtype)
+    minval = 2
+    maxval = 33
+    size = 1000
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=1234, alg=alg)
+      x = gen.uniform(
+          shape=[size], dtype=dtype, minval=minval, maxval=maxval).numpy()
+      self.assertTrue(np.all(x >= minval))
+      self.assertTrue(np.all(x <= maxval))
+
+  @parameterized.parameters(list(itertools.product(ALGS, FLOATS)))
+  def testNormalIsFinite(self, alg, dtype):
+    self.check_dtype(dtype)
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=1234, alg=alg)
+      x = gen.normal(shape=[10000], dtype=dtype).numpy()
+      self.assertTrue(np.all(np.isfinite(x)))
+
+  @parameterized.parameters(list(itertools.product(
+      ALGS, INTS + FLOATS, (12, 13, 123, 4321))))
+  def testDistributionOfUniform(self, alg, dtype, seed):
+    """Use Pearson's Chi-squared test to test for uniformity."""
+    self.check_dtype(dtype)
+    three_fry = random_ops_util.Algorithm.THREEFRY.value
+    auto_select = random_ops_util.Algorithm.AUTO_SELECT.value
+    is_tpu = xla_device().device_type == "TPU"
+    is_megacore = "megacore" in os.environ.get("TEST_TARGET", "").lower()
+    # TODO(b/244649364): Investigate why these combinations fail.
+    if ((alg, dtype, seed) in [(three_fry, dtypes.int64, 123),
+                               (three_fry, dtypes.uint64, 123)] or
+        is_tpu and
+        (alg, dtype, seed) in [(auto_select, dtypes.int64, 123),
+                               (auto_select, dtypes.uint64, 123)] or
+        is_megacore and
+        (alg, dtype, seed) in [(auto_select, dtypes.int32, 123),
+                               (auto_select, dtypes.uint32, 123),
+                               (auto_select, dtypes.int32, 12),
+                               (auto_select, dtypes.uint32, 12)]):
+      self.skipTest(
+          "This (device, alg, dtype, seed) combination fails (b/244649364).")
+    with ops.device(xla_device_name()):
+      n = 1000
+      gen = random.Generator.from_seed(seed=seed, alg=alg)
+      maxval = 1
+      if dtype.is_integer:
+        maxval = 100
+      t = gen.uniform(shape=[n], maxval=maxval, dtype=dtype)
+      x = t.numpy().astype(float)
+      if maxval > 1:
+        # Normalize y to range [0, 1).
+        x = x / maxval
+      # Tests that the values are distributed amongst 10 bins with equal
+      # probability. 16.92 is the Chi^2 value for 9 degrees of freedom with
+      # p=0.05. This test is probabilistic and would be flaky if the random
+      # seed were not fixed.
+      val = random_test_util.chi_squared(x, 10)
+      self.assertLess(val, 16.92)
+
+  @parameterized.parameters(list(itertools.product(ALGS, FLOATS)))
+  def testDistributionOfNormal(self, alg, dtype):
+    """Use Anderson-Darling test to test distribution appears normal."""
+    self.check_dtype(dtype)
+    with ops.device(xla_device_name()):
+      n = 1000
+      gen = random.Generator.from_seed(seed=1234, alg=alg)
+      x = gen.normal(shape=[n], dtype=dtype).numpy()
+      # The constant 2.492 is the 5% critical value for the Anderson-Darling
+      # test where the mean and variance are known. This test is probabilistic
+      # so to avoid flakiness the seed is fixed.
+      self.assertLess(
+          random_test_util.anderson_darling(x.astype(float)), 2.492)
+
+  @parameterized.parameters(list(itertools.product(ALGS, FLOATS)))
+  def testTruncatedNormal(self, alg, dtype):
+    self.check_dtype(dtype)
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=123, alg=alg)
+      n = 100000
+      y = gen.truncated_normal(shape=[n], dtype=dtype).numpy()
+      random_test_util.test_truncated_normal(
+          self.assertEqual, self.assertAllClose, n, y,
+          mean_atol=2e-3, median_atol=4e-3,
+          variance_rtol=1e-2 if dtype == dtypes.bfloat16 else 5e-3)
+
+  @test_util.disable_mlir_bridge(
+      "b/180412086: MLIR bridge gives wrong error messages.")
+  def testErrors(self):
+    """Tests that proper errors are raised.
+    """
+    shape = [2, 3]
+    with ops.device(xla_device_name()):
+      gen = random.Generator.from_seed(seed=1234, alg=random.RNG_ALG_THREEFRY)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          r"algorithm.* must be of shape \[\], not"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            gen.state.handle, [0, 0], shape)
+      with self.assertRaisesWithPredicateMatch(
+          TypeError, "EagerTensor of dtype int64"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            gen.state.handle, 1.1, shape)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          "Unsupported algorithm id"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            gen.state.handle, 123, shape)
+      with self.assertRaisesWithPredicateMatch(errors_impl.InvalidArgumentError,
+                                               "Unsupported algorithm id"):
+        gen_stateful_random_ops.rng_read_and_skip(
+            gen.state.handle, alg=123, delta=10)
+      var = variables.Variable([0, 0], dtype=dtypes.uint32)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          "Trying to read variable .* Expected int64 got"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            var.handle, random.RNG_ALG_THREEFRY, shape)
+      var = variables.Variable([[0]], dtype=dtypes.int64)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          "RNG state must have one and only one dimension, not"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            var.handle, random.RNG_ALG_THREEFRY, shape)
+      var = variables.Variable([0], dtype=dtypes.int64)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          "The size of the state must be at least"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            var.handle, random.RNG_ALG_THREEFRY, shape)
+      var = variables.Variable([0, 0], dtype=dtypes.int64)
+      with self.assertRaisesWithPredicateMatch(
+          errors_impl.InvalidArgumentError,
+          "The size of the state must be at least"):
+        gen_stateful_random_ops.stateful_standard_normal_v2(
+            var.handle, random.RNG_ALG_PHILOX, shape)
+
+
+if __name__ == "__main__":
+  ops.enable_eager_execution()
+  config.set_soft_device_placement(False)
+  test.main()