code: def testInteractiveIterationAxis2(self):
      def kernel(x):
          return x + 1, x
  
      def inner(x):
          state, trace = fun_mc.trace(x, kernel, 5)
          trace = jnp.transpose(trace, [1, 0])
          return state, trace
  
      state, trace = prefab.interactive_trace(
          jnp.zeros(2),
          inner,
          20,
          iteration_axis=2,
          progress_bar_fn=None,
      )
  
      self.assertAllClose([100.0, 100.0], state)
      self.assertEqual([2, 100], list(trace.shape))
      self.assertAllClose([99.0, 99.0], trace[:, -1])
  
imports: import functools
  import os
  import jax
  from jax import config
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import prefab
  from fun_mc import test_util
  from fun_mc import util_tfp
label_desc: Use image node_id 8 for calling the PrefabTest obj's underlying member method code with example usage: obj.testInteractiveIterationAxis2() and returns: unknown,  x, state,  trace