chore: import upstream snapshot with attribution

tensorflow/python/compiler/xla/tests/jit_compile_test.py

@@ -0,0 +1,101 @@
+# 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.
+# ==============================================================================
+
+from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
+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.ops import string_ops
+from tensorflow.python.platform import test
+
+
+class JitCompileTest(test.TestCase):
+
+  def testBasic(self):
+    with ops.Graph().as_default() as g:
+
+      def fn(x, a):
+        return x + a
+
+      xla_func = def_function.function(fn, jit_compile=True)
+      inputs = array_ops.placeholder(dtypes.float32, [5])
+      x = xla_func(inputs, 1)
+      with session.Session(graph=g) as sess:
+        y = sess.run(x, feed_dict={inputs: [1, 2, 2, 3, 3]})
+        self.assertTrue(x.graph.as_graph_def().library.function[0]
+                        .attr["_XlaMustCompile"].b)
+        self.assertAllClose([2, 3, 3, 4, 4], y)
+
+  def testDerivative(self):
+    def fn(x, a):
+      return 2 * x + a
+
+    with ops.Graph().as_default() as g:
+      xla_func = def_function.function(fn, jit_compile=True)
+      with backprop.GradientTape() as tape:
+        inputs = array_ops.placeholder(dtypes.float32, [5])
+        tape.watch(inputs)
+        outputs = xla_func(inputs, 1)
+      grads = tape.gradient(outputs, inputs)
+
+    with session.Session(graph=g) as sess:
+      grads_tensor = sess.run(grads, feed_dict={inputs: [1, 2, 2, 3, 3]})
+      self.assertAllClose([2, 2, 2, 2, 2], grads_tensor)
+      (forward, backward) = xla_func.get_concrete_function(
+          inputs, 1)._delayed_rewrite_functions.forward_backward()
+
+      # Check that the must-compile attribute gets correctly propagated to the
+      # created derivatives.
+      self.assertTrue(forward.cached_definition.attr["_XlaMustCompile"])
+      self.assertTrue(backward.function_def.attr["_XlaMustCompile"])
+
+  def testBasicInt32(self):
+    with ops.Graph().as_default() as g:
+
+      def fn(x, a):
+        return x + a
+
+      xla_func = def_function.function(fn, jit_compile=True)
+      inputs = array_ops.placeholder(dtypes.int32, [5])
+      x = xla_func(inputs, 1)
+      with session.Session(graph=g) as sess:
+        y = sess.run(x, feed_dict={inputs: [1, 2, 2, 3, 3]})
+        self.assertTrue(x.graph.as_graph_def().library.function[0]
+                        .attr["_XlaMustCompile"].b)
+        self.assertAllClose([2, 3, 3, 4, 4], y)
+
+  # Checking that we crash on an unsupported operation lets us test that the XLA
+  # compiler was actually invoked.
+  def testUnsupportedOps(self):
+    with ops.Graph().as_default() as g:
+
+      def fn(x):
+        return string_ops.string_length(
+            string_ops.string_format('{}', x))
+
+      xla_func = def_function.function(fn, jit_compile=True)
+      inputs = array_ops.placeholder(dtypes.float32, [5])
+      x = xla_func(inputs)
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  "Detected unsupported operations"):
+        with session.Session(graph=g) as sess:
+          sess.run(x, feed_dict={inputs: [1, 2, 2, 3, 3]})
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/jit_test.py

@@ -0,0 +1,315 @@
+# 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 python.compiler.xla.jit."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.compiler.xla import jit
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import function
+from tensorflow.python.framework import op_def_registry
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import random_seed
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gradients
+from tensorflow.python.ops import init_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+def enable_jit_nonstateful(node_def):
+  op_def = op_def_registry.get(node_def.op)
+  if op_def is None:
+    raise ValueError("Unregistered op being created: %s" % node_def)
+
+  return not op_def.is_stateful
+
+
+class JITTest(test.TestCase, parameterized.TestCase):
+
+  def compute(self, use_jit, compute_fn):
+    random_seed.set_random_seed(1234)
+    with self.session(graph=ops.Graph()) as sess:
+      with jit.experimental_jit_scope(use_jit):
+        r = compute_fn()
+      sess.run(variables.global_variables_initializer())
+      return (r, sess.run(r))
+
+  @test_util.run_v2_only
+  def testJITInEager(self):
+
+    with self.assertRaisesRegex(
+        RuntimeError, "xla.experimental.jit_scope is not supported when eager "
+        "execution is enabled. Try use it inside tf.function."):
+      with jit.experimental_jit_scope(True):
+        constant_op.constant(1)
+
+  @test_util.build_as_function_and_v1_graph
+  def testJITCreateOpsLambda(self):
+    """Test several ways of customizing the compilation attribute."""
+    def create_ops():
+      with variable_scope.variable_scope(
+          "root",
+          initializer=init_ops.random_uniform_initializer(
+              -0.1, 0.1, seed=2)):
+        inputs = random_ops.random_uniform((1,), minval=-10, maxval=10, seed=1)
+        return inputs
+    v_false_1_t, v_false_1 = self.compute(False, create_ops)
+    _, v_false_2 = self.compute(False, create_ops)
+    v_true_1_t, v_true_1 = self.compute(enable_jit_nonstateful, create_ops)
+    _, v_true_2 = self.compute(enable_jit_nonstateful, create_ops)
+    v_all_true_t, _ = self.compute(True, create_ops)
+    self.assertFalse(v_false_1_t.op.get_attr("_XlaCompile"))
+    v_true_1_t_sampler_op = v_true_1_t.graph.get_operation_by_name(
+        "root/random_uniform/RandomUniform")
+    v_all_true_t_sampler_op = v_all_true_t.graph.get_operation_by_name(
+        "root/random_uniform/RandomUniform")
+
+    self.assertFalse(v_true_1_t_sampler_op.get_attr("_XlaCompile"))
+    self.assertTrue(v_all_true_t_sampler_op.get_attr("_XlaCompile"))
+
+    self.assertTrue(v_true_1_t.op.get_attr("_XlaCompile"))
+    self.assertTrue(v_all_true_t.op.get_attr("_XlaCompile"))
+
+    # Additionally ensure that where no JIT compilation happens on the
+    # random_uniform op, the output values are identical to the case
+    # where no JIT compilation happens anywhere.
+    self.assertAllClose(v_false_1, v_false_2)
+    self.assertAllClose(v_true_1, v_true_2)
+    self.assertAllClose(v_false_1, v_true_1)
+
+  @test_util.build_as_function_and_v1_graph
+  def testJITXlaScope(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope(True):
+        # XlaScope 0
+        a1 = constant_op.constant(1)
+      with jit.experimental_jit_scope(True):
+        # XlaScope 1
+        a2 = constant_op.constant(1)
+        with jit.experimental_jit_scope(True):
+          # XlaScope still 1, depth 1
+          a3 = constant_op.constant(1)
+          with jit.experimental_jit_scope(True):
+            # XlaScope still 1, depth 2
+            a4 = constant_op.constant(1)
+          # XlaScope still 1, depth 1
+          a5 = constant_op.constant(1)
+      with jit.experimental_jit_scope(True):
+        # XlaScope now 2, depth 0
+        a6 = constant_op.constant(1)
+
+    self.assertEqual(b"jit_scope_0", a1.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a2.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a3.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a4.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_1", a5.op.get_attr("_XlaScope"))
+    self.assertEqual(b"jit_scope_2", a6.op.get_attr("_XlaScope"))
+
+  @test_util.build_as_function_and_v1_graph
+  def testJITVariableSeed(self):
+    """Test that the stateful initializer is not marked for compilation.
+
+    XLA does not currently support seeded initialization and XLA initializers
+    therefore return different values than non-XLA counterparts.  Here
+    we ensure that if we can disable JIT compilation for the initializers and
+    get the same variable values as if no JIT compilation happened.
+    """
+    def create_ops():
+      with variable_scope.variable_scope(
+          "root",
+          initializer=init_ops.random_uniform_initializer(
+              -0.1, 0.1, seed=2)):
+        inputs = variable_scope.get_variable("var", (1,))
+        return inputs
+    _, v_false_1 = self.compute(False, create_ops)
+    _, v_false_2 = self.compute(False, create_ops)
+    _, v_true_1 = self.compute(enable_jit_nonstateful, create_ops)
+    _, v_true_2 = self.compute(enable_jit_nonstateful, create_ops)
+    self.assertAllClose(v_false_1, v_false_2)
+    self.assertAllClose(v_true_1, v_true_2)
+    self.assertAllClose(v_false_1, v_true_1)
+
+  @test_util.build_as_function_and_v1_graph
+  def testDefunNoJitScope(self):
+    with self.session(graph=ops.Graph()):
+
+      @function.Defun(compiled=True, noinline=True)
+      def mulop(x1, x2):
+        return x1 * x2
+      x = constant_op.constant(1.0)
+      r = mulop(x, x)
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      # No enclosing jit scope so function sets its own value for _XlaScope.
+      self.assertEqual(b"function_mulop", func_attrs["_XlaScope"].s)
+
+  @test_util.build_as_function_and_v1_graph
+  def testDefunInheritsJitScope(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope(True):
+        @function.Defun(compiled=True, noinline=True)
+        def mulop(x1, x2):
+          return x1 * x2
+        x = constant_op.constant(1.0)
+        r = mulop(x, x)
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      # Ensure _XlaScope is inherited from enclosing context.
+      self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s)
+
+
+class CompilationEnabledInGradientTest(test.TestCase, parameterized.TestCase):
+
+  @test_util.build_as_function_and_v1_graph
+  def testCompilationInGradient(self):
+    with self.cached_session():
+      x = constant_op.constant([[3.]])
+      y_nc = math_ops.matmul(x, x, name="not_compiled")
+      with jit.experimental_jit_scope():
+        y_c = math_ops.matmul(y_nc, y_nc, name="compiled")
+      x_grads = gradients.gradients([y_c], [x])[0]
+      operations = x.graph.get_operations()
+      c_grad_ops = [
+          op for op in operations if "gradients/compiled" in op.name]
+      nc_grad_ops = [
+          op for op in operations if "gradients/not_compiled" in op.name]
+      self.assertGreater(len(c_grad_ops), 0)
+      self.assertGreater(len(nc_grad_ops), 0)
+      for cg in c_grad_ops:
+        self.assertTrue(cg.get_attr("_XlaCompile"))
+      for ncg in nc_grad_ops:
+        with self.assertRaisesRegex(ValueError, "[Nn]o attr named"):
+          ncg.get_attr("_XlaCompile")
+
+      # d/dx (x ** 4) = 4 * (x ** 3)
+      self.assertAllClose([[108]], x_grads)
+
+  @test_util.build_as_function_and_v1_graph
+  def testCompilationGradientScopeNames(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope():
+        # XlaScope 0
+        a1 = constant_op.constant([[1.]])
+        a1t = math_ops.matmul(a1, a1)
+      with jit.experimental_jit_scope():
+        # XlaScope 1
+        a2 = constant_op.constant([[1.]])
+        a2t = math_ops.matmul(a2, a2)
+
+      self.assertEqual(b"jit_scope_0", a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1", a2.op.get_attr("_XlaScope"))
+      grad_a1 = gradients.gradients(a1t, a1, name="GA")[0]
+      grad_a2 = gradients.gradients(a2t, a2, name="GB")[0]
+      grad_a1 = grad_a1.op.inputs[0]
+      grad_a2 = grad_a2.op.inputs[0]
+      self.assertTrue(grad_a1.op.get_attr("_XlaCompile"))
+      self.assertTrue(grad_a2.op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0", grad_a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1", grad_a2.op.get_attr("_XlaScope"))
+
+  @test_util.build_as_function_and_v1_graph
+  def testCompilationSeparateGradientScopeNames(self):
+    with self.session(graph=ops.Graph()):
+      with jit.experimental_jit_scope(True, separate_compiled_gradients=True):
+        # XlaScope 0
+        a1 = constant_op.constant([[1.]])
+        a1t = math_ops.matmul(a1, a1)
+      with jit.experimental_jit_scope(True, separate_compiled_gradients=True):
+        # XlaScope 1
+        a2 = constant_op.constant([[1.]])
+        a2t = math_ops.matmul(a2, a2)
+
+      self.assertEqual(b"jit_scope_0", a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1", a2.op.get_attr("_XlaScope"))
+      grad_a1 = gradients.gradients(a1t, a1, name="GA")[0]
+      grad_a2 = gradients.gradients(a2t, a2, name="GB")[0]
+      grad_a1 = grad_a1.op.inputs[0]
+      grad_a2 = grad_a2.op.inputs[0]
+      self.assertTrue(grad_a1.op.get_attr("_XlaCompile"))
+      self.assertTrue(grad_a2.op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0_grad_GA",
+                       grad_a1.op.get_attr("_XlaScope"))
+      self.assertEqual(b"jit_scope_1_grad_GB",
+                       grad_a2.op.get_attr("_XlaScope"))
+
+  @test_util.build_as_function_and_v1_graph
+  def testPlaysNicelyWithDefun(self):
+    with self.session(graph=ops.Graph()) as sess:
+      with jit.experimental_jit_scope(True):
+        @function.Defun(compiled=True, noinline=True)
+        def mulop(x1, x2):
+          return x1 * x2
+        x = constant_op.constant(1.0)
+        r = mulop(x, x)
+        g_r = gradients.gradients(r, x, name="GA")[0]
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s)
+
+      # Ensure the gradient (SymbolicGradient) is compiled, with the same
+      # _XlaScope as the function itself.
+      grad_op = g_r.op.inputs[0].op
+      self.assertTrue(grad_op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0", grad_op.get_attr("_XlaScope"))
+
+      # Ensure the ops run: grad(x1*x1) = 2*x1
+      self.assertAllClose([1.0, 1.0, 2.0], sess.run([x, r, g_r]))
+
+  @test_util.build_as_function_and_v1_graph
+  def testPlaysNicelyWithDefunSeparateGradientScope(self):
+    with self.session(graph=ops.Graph()) as sess:
+      with jit.experimental_jit_scope(True):
+
+        @function.Defun(
+            compiled=True, noinline=True, separate_compiled_gradients=True)
+        def mulop(x1, x2):
+          return x1 * x2
+
+        x = constant_op.constant(1.0)
+        r = mulop(x, x)
+        g_r = gradients.gradients(r, x, name="GA")[0]
+
+      # Ensure the forward function is compiled.
+      graph_def = r.graph.as_graph_def()
+      func_attrs = graph_def.library.function[0].attr
+      self.assertTrue(func_attrs["_XlaCompile"].b)
+      self.assertEqual(b"jit_scope_0", func_attrs["_XlaScope"].s)
+
+      # Ensure the gradient (SymbolicGradient) is compiled, with a different
+      # _XlaScope from the function itself.
+      grad_op = g_r.op.inputs[0].op
+      self.assertTrue(grad_op.get_attr("_XlaCompile"))
+      self.assertEqual(b"jit_scope_0_grad_GA",
+                       grad_op.get_attr("_XlaScope"))
+
+      # Ensure the ops run: grad(x1*x1) = 2*x1
+      self.assertAllClose([1.0, 1.0, 2.0], sess.run([x, r, g_r]))
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/pjrt_autoclustering_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 single device compilation + autoclustering using the Device API (PjRt).
+
+This feature is still under active development and is protected behind the
+`--tf_xla_use_device_api` flag in the `TF_XLA_FLAGS` environment variable.
+"""
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import math_ops
+
+
+class PjrtAutoclusteringTest(test.TestCase):
+
+  def test_xla_compile_and_run_on_gpu_device(self):
+
+    if not test.is_gpu_available() or not test.is_built_with_gpu_support():
+      test.skipTest("Test only applicable on GPU")
+
+    @def_function.function
+    def arithmetic(x):
+      return 2 * x + 1
+
+    @def_function.function
+    def conditional(x):
+      # cond uses switch and merge, which are not supported by XLA based on
+      # https://docs.google.com/spreadsheets/d/1H8AIDdnlyyaWZOYN3WpBVNOmGOS_M8OyF7IA7kL3fjk/edit?resourcekey=0-I-mIp472YuK8FuBa5Zmzmg#gid=139369773
+      return cond.cond(math_ops.reduce_sum(x) < 5, lambda: x + x, lambda: x)
+
+    @def_function.function
+    def func(x, y):
+      return (arithmetic(x) + conditional(y) ** 2) / 2
+
+    i1 = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+
+    # Simple case: all ops supported by XLA
+    with ops.device("/device:GPU:0"):
+      with context.collect_graphs(optimized=True) as graphs:
+        result = arithmetic(i1)
+    self.assertAllClose(result.numpy(), [[3.0, 5.0], [7.0, 9.0]], atol=1e-05)
+    graph_ops = [n.op for n in graphs[0].node]
+    self.assertContainsSubset(["_XlaCompile", "_XlaRun"], graph_ops)
+
+    # Complex case: includes ops not supported by XLA (switch and merge)
+    i2 = constant_op.constant([[5.0, 6.0], [7.0, 8.0]])
+    with ops.device("/device:GPU:0"):
+      with context.collect_graphs(optimized=True) as graphs:
+        result = func(i1, i2)
+    self.assertAllClose(result.numpy(), [[14.0, 20.5], [28, 36.5]], atol=1e-05)
+    graph_ops = [n.op for n in graphs[0].node]
+    self.assertContainsSubset(["_XlaCompile", "_XlaRun"], graph_ops)
+    # because of the cond, not all ops can be combined into a single _XlaCompile
+    self.assertGreater(graph_ops.count("_XlaCompile"), 1)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/pjrt_compile_test.py

@@ -0,0 +1,162 @@
+# 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 single device compilation + execution using the Device API (aka PjRt).
+
+This feature is still under active development and is protected behind the
+`--tf_xla_use_device_api` flag in the `TF_XLA_FLAGS` environment variable.
+"""
+import numpy as np
+from tensorflow.python.eager import def_function
+from tensorflow.python.eager import test
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import variables
+
+
+class PjrtCompileTest(test.TestCase):
+
+  def test_compile_on_demand(self):
+    if not test.is_gpu_available() or not test.is_built_with_gpu_support():
+      test.skipTest("Test only applicable on GPU")
+
+    with ops.device("/device:XLA_GPU:0"):
+      a = constant_op.constant([1.0, 2.0])
+      b = constant_op.constant([2.0, 3.0])
+      c = a + b
+      self.assertAllClose([3.0, 5.0], c, atol=1e-05)
+
+      v = variables.Variable([0.0, 1.0])
+      v.assign([1.0, 2.0])
+      self.assertAllClose([1.0, 2.0], v.value(), atol=1e-05)
+      v.assign_add([1.0, 2.0])
+      self.assertAllClose([2.0, 4.0], v.value(), atol=1e-05)
+
+      d = c + v
+      self.assertAllClose([5.0, 9.0], d, atol=1e-05)
+
+  # Tests compilation and execution of a jit_compiled function using PjRt.
+  def test_xla_local_launch(self):
+    if not test.is_gpu_available() or not test.is_built_with_gpu_support():
+      test.skipTest("Test only applicable on GPU")
+
+    @def_function.function(jit_compile=True)
+    def foo(x, y):
+      return x + y + 1
+
+    @def_function.function(jit_compile=True)
+    def bar(x, y):
+      x.assign(y)
+      y.assign_add([1.0, 1.0])
+
+    # Currently PjRt only supports compilation and execution for the XLA_GPU
+    # device to unblock development. Support for non-XLA devices (CPU/GPU/single
+    # core TPU) is going to be added soon, after which support for XLA_* devices
+    # will be dropped.
+    # TODO(b/255826209): Modify the test as we progress towards supporting
+    # non-XLA devices.
+    with ops.device("/device:XLA_GPU:0"):
+      # Function call with scalars
+      self.assertEqual(self.evaluate(foo(1, 2)), 4)
+
+      # Function call with tensors
+      a = constant_op.constant([1.0, 2.0])
+      b = constant_op.constant([2.0, 3.0])
+      self.assertAllClose([4.0, 6.0], foo(a, b), atol=1e-05)
+
+      # Function call with variables
+      x = variables.Variable([0.0, 1.0])
+      y = variables.Variable([1.0, 2.0])
+      self.assertAllClose([2.0, 4.0], foo(x, y), atol=1e-05)
+
+      # Function call with constant and variable
+      self.assertAllClose([2.0, 4.0], foo(a, x), atol=1e-05)
+
+      # Function call that updates variables
+      bar(x, y)
+      self.assertAllClose([1.0, 2.0], x.value(), atol=1e-05)
+      self.assertAllClose([2.0, 3.0], y.value(), atol=1e-05)
+
+  def test_xla_launch_and_tf_kernel_on_gpu_device(self):
+    @def_function.function(jit_compile=True)
+    def const_fn():
+      return constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+
+    @def_function.function(jit_compile=True)
+    def matmul_fn(x):
+      return math_ops.matmul(x, x)
+
+    # The following tests XlaLaunch kernel interoperation with legacy TF
+    # GPU kernel.
+    #
+    # The following use case is tested:
+    # host -> h2d transfer (creating a tf::Tensor with PjRtTensorBuffer)
+    # -> XLA op (GPU) -> TF op (GPU) -> d2h transfer.
+    host_tensor = constant_op.constant([[1.0, 2.0], [3.0, 4.0]])
+    with ops.device("/device:GPU:0"):
+      xla_tensor = const_fn()
+      xla_result = matmul_fn(host_tensor)
+      result = math_ops.matmul(xla_result, xla_tensor)  # TF GPU kernel.
+
+    # Using numpy to calculate the reference result ("ref_*").
+    ref_tensor = np.array([[1.0, 2.0], [3.0, 4.0]])
+    ref_result = np.matmul(np.matmul(ref_tensor, ref_tensor), ref_tensor)
+    self.assertAllClose(result.numpy(), ref_result, atol=1e-05)
+
+    # The following use case is tested:
+    # host -> h2d transfer (creating a tf::Tensor with PjRtTensorBuffer)
+    # -> TF op (GPU, creating a tf::Tensor with device mem)
+    # -> XLA op (GPU, accepting a tf::Tensor with device mem) -> d2h transfer.
+    with ops.device("/device:GPU:0"):
+      tf_matmul_tensor = math_ops.matmul(host_tensor, host_tensor)
+      xla_result = matmul_fn(tf_matmul_tensor)
+
+    ref_matmul_tensor = np.matmul(ref_tensor, ref_tensor)
+    ref_result_2 = np.matmul(ref_matmul_tensor, ref_matmul_tensor)
+    self.assertAllClose(xla_result.numpy(), ref_result_2, atol=1e-05)
+
+  def test_xla_launch_with_var_on_gpu_device(self):
+    @def_function.function(jit_compile=True)
+    def foo(x, y):
+      return x + y + 1
+
+    @def_function.function(jit_compile=True)
+    def bar(x, y):
+      x.assign(y)
+      y.assign_add([1.0, 1.0])
+
+    with ops.device("/device:GPU:0"):
+      a = constant_op.constant([1.0, 2.0])
+      x = variables.Variable([0.0, 1.0])
+      result_tensor = foo(x, a)
+
+    self.assertAllClose(result_tensor.numpy(), [2.0, 4.0], atol=1e-05)
+
+    # The following use case is tested:
+    # Variable updates following an XLA computation that reads the updated
+    # variables.
+    with ops.device("/device:GPU:0"):
+      var_a = variables.Variable([0.0, 1.0])
+      var_b = variables.Variable([1.0, 2.0])
+      bar(var_a, var_b)
+      result = foo(var_a, var_b)
+
+    self.assertAllClose([1.0, 2.0], var_a.value(), atol=1e-05)
+    self.assertAllClose([2.0, 3.0], var_b.value(), atol=1e-05)
+    self.assertAllClose(result, [4.0, 6.0], atol=1e-05)
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/compiler/xla/tests/xla_test.py

@@ -0,0 +1,337 @@
+# 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 python.compiler.xla.xla."""
+
+from absl.testing import parameterized
+
+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 ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import control_flow_util
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import state_ops
+from tensorflow.python.ops import variable_scope
+from tensorflow.python.ops import while_loop
+from tensorflow.python.platform import test
+from tensorflow.python.summary import summary
+from tensorflow.python.tpu import tpu_feed
+
+
+_EXPECTED_LOSS = 1
+_EXPECTED_FEATURE = 2
+_EXPECTED_LABEL = 3
+
+
+class XLACompileContextTest(test.TestCase, parameterized.TestCase):
+
+  def create_test_xla_compile_context(self):
+    computation_name = ops.get_default_graph().unique_name('computation')
+    pivot = control_flow_ops.no_op(name=computation_name + '/pivot')
+    return xla.XLACompileContext(name=computation_name, pivot=pivot)
+
+  @test_util.run_v1_only('Testing graph mode behavior only')
+  def test_report_unsupported_operations_graph_mode(self):
+    """Tests that unsupported operations are detected."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    dummy_tensor = constant_op.constant(1.1)
+    audio_summary = summary.audio('audio_summary', dummy_tensor, 0.5)
+    histogram_summary = summary.histogram('histogram_summary', dummy_tensor)
+    image_summary = summary.image('image_summary', dummy_tensor)
+    scalar_summary = summary.scalar('scalar_summary', dummy_tensor)
+    tensor_summary = summary.tensor_summary('tensor_summary', dummy_tensor)
+    summary.merge(
+        [
+            audio_summary, histogram_summary, image_summary, scalar_summary,
+            tensor_summary
+        ],
+        name='merge_summary')
+    logging_ops.Print(dummy_tensor, [dummy_tensor], name='print_op')
+    context.Exit()
+
+    unsupported_ops_names = [op.name for op in context._unsupported_ops]
+    self.assertEqual(unsupported_ops_names, [
+        u'audio_summary', u'histogram_summary', u'image_summary',
+        u'scalar_summary', u'tensor_summary', u'merge_summary/merge_summary',
+        u'print_op'
+    ])
+
+  @test_util.run_v1_only('Testing graph mode behavior only')
+  def test_resource_variable_graph_mode(self):
+    """Tests that resource variable usage is allowed."""
+    a = variable_scope.get_variable(
+        name='variable_a', use_resource=True, initializer=1)
+
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    a.assign(2)
+    context.Exit()
+
+  def test_resource_variable_in_function(self):
+    """Tests that resource variable usage is allowed."""
+    a = variable_scope.get_variable(
+        name='variable_a', use_resource=True, initializer=1)
+
+    @def_function.function
+    def func():
+      context = self.create_test_xla_compile_context()
+      context.Enter()
+      o = a.assign(2)
+      context.Exit()
+      return o
+
+    self.assertEqual(self.evaluate(func()), 2)
+
+  @test_util.run_v1_only('Testing v1-only ref variable handling.')
+  def test_non_resource_variable_error(self):
+    """Tests that non-resource variable usage is disallowed."""
+    a = variable_scope.get_variable(
+        name='variable_a', shape=(1), use_resource=False)
+
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    with self.assertRaisesRegex(
+        NotImplementedError, 'Non-resource Variables are not supported inside '
+        r'XLA computations \(operator name: Assign\)'):
+      state_ops.assign(a, a + 1)
+    context.Exit()
+
+  @test_util.build_as_function_and_v1_graph
+  def test_nested_xla_compile_error(self):
+    """Tests that nested XLA computation leads to fatal error."""
+    context1 = self.create_test_xla_compile_context()
+    context1.Enter()
+
+    context2 = self.create_test_xla_compile_context()
+    context2.Enter()
+    with self.assertRaisesRegex(ValueError,
+                                'XLA compiled computations cannot be nested'):
+      constant_op.constant(1)
+    context2.Exit()
+    context1.Exit()
+
+  @test_util.build_as_function_and_v1_graph
+  def test_xla_compile_attr(self):
+    """Tests that ops are tagged with XLA compile ID attribute."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+    self.assertIn('_xla_compile_id', op.op.node_def.attr)
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_without_input(self):
+    """Tests that ops without inputs depend on pivot correctly."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+
+    self.assertIn(context._pivot, op.op.control_inputs)
+
+  @test_util.run_v1_only('Testing graph mode behavior only')
+  def test_external_control_edges_graph_mode(self):
+    """Tests that external control edges are handled correctly."""
+    i = constant_op.constant(1)
+    op1 = constant_op.constant(1)
+
+    with ops.control_dependencies([op1]):
+      op2 = constant_op.constant(1)
+    self.assertIn(op1.op, op2.op.control_inputs)
+
+    def while_body(i):
+      del i  # unused
+      context = self.create_test_xla_compile_context()
+      context.Enter()
+      with ops.control_dependencies([op1]):
+        op3 = constant_op.constant(1)
+      context.Exit()
+      self.assertNotIn(op1.op, op3.op.control_inputs)
+      return op3
+
+    while_loop.while_loop(
+        cond=lambda i: math_ops.less(i, 10), body=while_body, loop_vars=[i])
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_output_marked_as_seen(self):
+    """Tests that any op output is marked as seen in context."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+
+    self.assertIn(op.name, context._values)
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_is_in_context(self):
+    """Tests that XLACompileContext is recognized as an XLA context."""
+    op1 = constant_op.constant(1)
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op2 = constant_op.constant(2)
+    context.Exit()
+    self.assertFalse(control_flow_util.IsInXLAContext(op1.op))
+    self.assertTrue(control_flow_util.IsInXLAContext(op2.op))
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_prevent_feeding(self):
+    """Tests that ops created inside XLACompileContext can not be fed."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+    self.assertFalse(op.graph.is_feedable(op.op))
+
+  @test_util.build_as_function_and_v1_graph
+  def test_op_prevent_fetching(self):
+    """Tests that ops created inside XLACompileContext can not be fetched."""
+    context = self.create_test_xla_compile_context()
+    context.Enter()
+    op = constant_op.constant(1)
+    context.Exit()
+    self.assertFalse(op.graph.is_fetchable(op.op))
+
+
+class XlaCompileTest(test.TestCase):
+
+  @test_util.run_v2_only
+  @test_util.disable_tfrt(
+      'Legacy XLA test. It depends on EncapsulateXlaComputationsPass.')
+  def test_xla_compile_eager(self):
+    """Tests that xla.compile raises proper exception when used eagerly."""
+
+    def computation(a, b):
+      return a + b
+
+    self.assertEqual(self.evaluate(xla.compile(computation, [1, 2])[0]), 3)
+
+  @test_util.disable_tfrt(
+      'Legacy XLA test. It depends on EncapsulateXlaComputationsPass.')
+  def test_xla_compile_in_function(self):
+    """Tests that xla.compile works in tf.function."""
+
+    @def_function.function
+    def func_wrapper(a):
+
+      def compute(a):
+        return a + 1
+
+      return xla.compile(compute, [a])
+
+    self.assertEqual(self.evaluate(func_wrapper(1))[0], 2)
+
+  @test_util.disable_tfrt(
+      'Legacy XLA test. It depends on EncapsulateXlaComputationsPass.')
+  def test_xla_compile_write_variable_in_function(self):
+    """Tests that xla.compile works with variable in tf.function."""
+    a = variable_scope.get_variable(
+        name='variable_a', use_resource=True, initializer=1)
+
+    @def_function.function
+    def func_wrapper():
+
+      def compute():
+        a.assign_add(1)
+        a.assign_sub(2)
+        return a.read_value()
+
+      return xla.compile(compute)
+
+    self.evaluate(a.initializer)
+    self.assertEqual(self.evaluate(func_wrapper())[0], 0)
+
+
+class CheckFunctionArgumentCountTest(test.TestCase):
+
+  def test_simple(self):
+    """Tests that arg checker works for functions with no varargs or defaults.
+    """
+
+    def func(x, y, z):
+      return x + y + z
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual('exactly 3 arguments',
+                     xla.check_function_argument_count(func, 2, None))
+    queue = tpu_feed.InfeedQueue(2)
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual('exactly 3 arguments',
+                     xla.check_function_argument_count(func, 2, queue))
+
+  def test_default_args(self):
+    """Tests that arg checker works for a function with no varargs."""
+
+    def func(x, y, z=17):
+      return x + y + z
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, None))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 1, None))
+    self.assertEqual('at most 3 arguments',
+                     xla.check_function_argument_count(func, 4, None))
+    queue = tpu_feed.InfeedQueue(1)
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 0, queue))
+    self.assertEqual('at most 3 arguments',
+                     xla.check_function_argument_count(func, 4, queue))
+
+  def test_var_args(self):
+    """Tests that arg checker works for a function with varargs."""
+
+    def func(x, y, *z):
+      return x + y + len(z)
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 4, None))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 1, None))
+    queue = tpu_feed.InfeedQueue(1)
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, queue))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 0, queue))
+
+  def test_var_args_and_defaults(self):
+    """Tests that arg checker works for a function with varargs and defaults."""
+
+    def func(x, y, z=17, *q):  # pylint: disable=keyword-arg-before-vararg
+      return x + y + z + len(q)
+
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 4, None))
+    self.assertEqual(None, xla.check_function_argument_count(func, 5, None))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 1, None))
+    queue = tpu_feed.InfeedQueue(1)
+    self.assertEqual(None, xla.check_function_argument_count(func, 1, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 2, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 3, queue))
+    self.assertEqual(None, xla.check_function_argument_count(func, 4, queue))
+    self.assertEqual('at least 2 arguments',
+                     xla.check_function_argument_count(func, 0, queue))
+
+
+if __name__ == '__main__':
+  test.main()