code: def test_known_size(self):
      tracer = tracing_reducer.TracingReducer(size=3)
      self.assertEqual(tracer.size, 3)
      state = tracer.initialize(tf.zeros(()), tf.zeros(()))
      for sample in range(3):
          state = tracer.one_step(sample, state, sample)
      all_states, final_trace = tracer.finalize(state)
      self.assertAllEqual(
          [3], tensorshape_util.as_list(all_states.shape)
      )
      self.assertAllEqual(
          [3], tensorshape_util.as_list(final_trace.shape)
      )
      all_states, final_trace = self.evaluate([all_states, final_trace])
      self.assertAllEqual([0, 1, 2], all_states)
      self.assertAllEqual([0, 1, 2], final_trace)
  
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 3 for calling the TracingReducerTest obj's underlying member method code with example usage: obj.test_known_size() without return types