code: def testCorrectlyWorksWithCorrection(self):
      current_state_ = [self.dtype([1, 2]), self.dtype([3, 4])]
      current_state = [tf.convert_to_tensor(s) for s in current_state_]
  
      expected_inner_init_kernel_results = (
          InnerKernelResultsWithCorrection(
              log_acceptance_correction=self.dtype([+300.0, -300.0]),
              target_log_prob=self.dtype([100.0, -100.0]),
              grads_target_log_prob=[
                  self.dtype([1.25, 1.5]),
                  self.dtype([2.25, 2.5]),
              ],
              extraneous=self.dtype([1.75, 2.0]),
          )
      )
  
      one_step_fn = make_one_step_fn(dtype=self.dtype)
      bootstrap_fn = make_bootstrap_results_fn(
          expected_inner_init_kernel_results
      )
  
      # Collect expected results.
      expected_init_inner_kernel_results = bootstrap_fn(current_state)
      _, expected_inner_kernel_results = one_step_fn(
          current_state, expected_init_inner_kernel_results
      )
  
      # Collect actual results.
      mh = metropolis_hastings.MetropolisHastings(
          FakeTransitionKernel(
              is_calibrated=False,
              one_step_fn=one_step_fn,
              bootstrap_fn=bootstrap_fn,
          )
      )
      init_kernel_results = mh.bootstrap_results(current_state)
      next_state, kernel_results = mh.one_step(
          current_state, init_kernel_results, seed=test_util.test_seed()
      )
  
      # Unmodified state is passed through unmodified.
      self.assertIs(
          kernel_results.accepted_results.extraneous,
          init_kernel_results.accepted_results.extraneous,
      )
      self.assertIs(
          kernel_results.proposed_results.extraneous,
          init_kernel_results.accepted_results.extraneous,
      )
  
      # Check correct types and call pattern.
      self.assertEqual(
          dict(is_calibrated=1, one_step=1, bootstrap_results=1),
          mh.inner_kernel.call_count,
      )
      for kr in [
          init_kernel_results.accepted_results,
          init_kernel_results.proposed_results,
      ]:
          self.assertEqual(
              type(expected_init_inner_kernel_results), type(kr)
          )
      for kr in [
          kernel_results.accepted_results,
          kernel_results.proposed_results,
      ]:
          self.assertEqual(
              type(expected_inner_kernel_results), type(kr)
          )
  
      # Now check actual values.
      [
          expected_init_inner_kernel_results_,
          expected_inner_kernel_results_,
          init_kernel_results_,
          kernel_results_,
          next_state_,
      ] = self.evaluate(
          [
              expected_init_inner_kernel_results,
              expected_inner_kernel_results,
              init_kernel_results,
              kernel_results,
              next_state,
          ]
      )
  
      # Check that the bootstrapped kernel results are correctly initialized.
      for fn in expected_inner_init_kernel_results._fields:
          self.assertAllClose(
              getattr(expected_init_inner_kernel_results_, fn, np.nan),
              getattr(
                  init_kernel_results_.accepted_results, fn, np.nan
              ),
              atol=0.0,
              rtol=1e-5,
          )
  
      # Check that the proposal is correctly computed.
      self.assertAllClose(
          [2 * current_state_[0], 3 * current_state_[1]],
          kernel_results_.proposed_state,
          atol=0.0,
          rtol=1e-5,
      )
      for fn in expected_inner_kernel_results._fields:
          self.assertAllClose(
              getattr(expected_inner_kernel_results_, fn, np.nan),
              getattr(kernel_results_.proposed_results, fn, np.nan),
              atol=0.0,
              rtol=1e-5,
          )
  
      # First: Extremely high correction means proposed will be accepted, despite
      #        high prob initial state.
      # Second: Extremely low correction means proposed will be rejected, despite
      #         low prob initial state.
      self.assertAllEqual([True, False], kernel_results_.is_accepted)
      self.assertAllEqual(
          [
              (0.5 * 100.0) - (100.0) + (0.5 * 300.0),
              (0.5 * -100.0) - (-100.0) + (0.5 * -300.0),
          ],
          kernel_results_.log_accept_ratio,
      )
      self.assertAllClose(
          [
              self.dtype([0 + 2, 1]) * current_state_[0],
              self.dtype([1 + 2, 1]) * current_state_[1],
          ],
          next_state_,
      )
  
imports: import collections
  import warnings
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc import kernel
  from tensorflow_probability.python.mcmc import metropolis_hastings
  from tensorflow_probability.python.mcmc.internal.util import is_list_like
label_desc: Use image node_id 3 for calling the MetropolisHastingsTest obj's underlying member method code with example usage: obj.testCorrectlyWorksWithCorrection() without return types