code: 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)
  
file_path: tensorflow-upstream/tensorflow/python/compiler/xla/tests/pjrt_compile_test.py
imports: import numpy
  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
label_desc: Use image node_id 2 for calling the PjrtCompileTest obj's underlying member method code with example usage: obj.test_xla_local_launch() and returns: unknown