code: def test_stateless_while(self):
      x = random_ops.random_uniform([3, 5])
      lengths = constant_op.constant([4, 0, 2])
  
      def loop_fn(i):
          x_i = array_ops.gather(x, i)
          lengths_i = array_ops.gather(lengths, i)
  
          return while_loop.while_loop(
              lambda j, _: j < lengths_i,
              lambda j, t: (j + 1, t + array_ops.gather(x_i, j)),
              [0, 0.0],
          )
  
      self._test_loop_fn(loop_fn, 3)
  
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 4 for calling the WhileV2Test obj's underlying member method code with example usage: obj.test_stateless_while() and returns: while_loop