code: 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)
  
file_path: tensorflow/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 3 for calling the PjrtCompileTest obj's underlying member method code with example usage: obj.test_xla_launch_and_tf_kernel_on_gpu_device() and returns: constant_op, math_ops