code: def test_tf_while(self):
      def trace_fn(sample, pkr):
          return sample, (sample, pkr), {"one": sample, "two": pkr}
  
      tracer = tracing_reducer.TracingReducer(trace_fn=trace_fn)
      state = tracer.initialize(tf.zeros(()), tf.zeros(()))
  
      def _body(sample, pkr, state):
          new_state = tracer.one_step(sample, state, pkr)
          return (sample + 1, pkr + 2, new_state)
  
      _, _, state = tf.while_loop(
          cond=lambda i, _, __: i < 3,
          body=_body,
          loop_vars=(1.0, 2.0, state),
      )
      final_trace = self.evaluate(tracer.finalize(state))
      self.assertEqual(3, len(final_trace))
      self.assertAllEqual([1, 2], final_trace[0])
      self.assertAllEqual(([1, 2], [2, 4]), final_trace[1])
      self.assertAllEqualNested(
          final_trace[2], ({"one": [1, 2], "two": [2, 4]})
      )
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/tracing_reducer_test.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.mcmc import tracing_reducer
  from tensorflow_probability.python.experimental.mcmc.internal import test_fixtures
  from tensorflow_probability.python.experimental.mcmc.sample_fold import sample_fold
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the TracingReducerTest obj's underlying member method code with example usage: obj.test_tf_while() and returns: sample,  sample,  pkr,  dict, unknown,  unknown,  new_state