class_docstr: Base class for test cases.
code: def _test_loop_fn(
      self,
      loop_fn,
      iters,
      parallel_iterations=None,
      fallback_to_while_loop=False,
      rtol=1e-4,
      atol=1e-5,
  ):
      t1 = pfor_control_flow_ops.pfor(
          loop_fn,
          iters=iters,
          fallback_to_while_loop=fallback_to_while_loop,
          parallel_iterations=parallel_iterations,
      )
      loop_fn_dtypes = nest.map_structure(lambda x: x.dtype, t1)
      t2 = pfor_control_flow_ops.for_loop(
          loop_fn,
          loop_fn_dtypes,
          iters=iters,
          parallel_iterations=parallel_iterations,
      )
  
      def _check_shape(a, b):
          msg = (
              "Inferred static shapes are different between two loops:"
              f" {a.shape} vs {b.shape}."
          )
          # TODO(b/268146947): should assert bool(a.shape) == bool(b.shape),
          # since both should be either defined or undefined. But it does not work.
          if b.shape:
              self.assertEqual(
                  a.shape.as_list()[0], b.shape.as_list()[0], msg
              )
              # TODO(b/268146947): self.assertShapeEqual(a, b, msg) does not work.
  
      nest.map_structure(_check_shape, t1, t2)
      self.run_and_assert_equal(t1, t2, rtol=rtol, atol=atol)
  
file_path: tensorflow-upstream/tensorflow/python/ops/parallel_for/test_util.py
imports: import numpy
  from tensorflow.python.ops import variables
  from tensorflow.python.ops.parallel_for import control_flow_ops
  from tensorflow.python.platform import test
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling the PForTestCase obj's underlying member method code with example usage: obj._test_loop_fn(loop_fn, iters, parallel_iterations, fallback_to_while_loop, rtol, atol) without return types