code: def test_xla_while_loop(self):
      def compute(x):
          return math_ops.reduce_mean(x, axis=0, keepdims=True)
  
      def vectorized_compute(x, i):
          inp = array_ops.gather(x, i)
          output = pfor_control_flow_ops.vectorized_map(compute, inp)
          output.set_shape([5, 1])
          return output
  
      def while_compute(x):
          return while_loop.while_loop_v2(
              lambda i, _: i < 10,
              lambda i, y: (i + 1, y + vectorized_compute(x, i)),
              (0, array_ops.zeros([5, 1])),
          )[1]
  
      result = xla.compile(
          while_compute, inputs=[array_ops.ones((10, 5, 3))]
      )
      expected = array_ops.ones([5, 1]) * 10
      self.run_and_assert_equal(expected, result)
  
imports: from tensorflow.compiler.tf2xla.python import xla
  from tensorflow.python.compiler.xla import jit
  from tensorflow.python.compiler.xla import xla
  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 test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_v2_toggles
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import random_ops
  from tensorflow.python.ops import resource_variable_ops
  from tensorflow.python.ops import while_loop
  from tensorflow.python.ops.parallel_for import control_flow_ops
  from tensorflow.python.ops.parallel_for.test_util import PForTestCase
  from tensorflow.python.platform import test
label_desc: Use image node_id 5 for calling the PForTest obj's underlying member method code with example usage: obj.test_xla_while_loop() and returns: math_ops, output, while_loop