code: def test_transform_fn_with_nested_return(self):
      cov_red = covariance_reducer.CovarianceReducer(
          transform_fn=lambda sample, _: [sample, sample + 1]
      )
      fake_kr = FakeKernelResults(0.0, FakeInnerResults(0))
      state = cov_red.initialize(tf.zeros((2,)), fake_kr)
      _, state = tf.while_loop(
          lambda sample, _: sample < 5,
          lambda s, st: (
              s + 1,
              cov_red.one_step(tf.ones((2,)) * s, st, fake_kr),
          ),
          (0.0, state),
      )
      final_cov = self.evaluate(cov_red.finalize(state))
      self.assertEqual(2, len(final_cov))
      self.assertAllEqual(np.ones((2, 2, 2)) * 2, final_cov)
  
file_path: probability/tensorflow_probability/python/experimental/mcmc/covariance_reducer_test.py
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.experimental.mcmc import covariance_reducer
  from tensorflow_probability.python.experimental.mcmc.internal import test_fixtures
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 6 for calling the CovarianceReducersTest obj's underlying member method code with example usage: obj.test_transform_fn_with_nested_return() without return types