code: def test_non_autobatched_joint_distribution(self):
      def transition_fn(_, previous_state):
          return jdn.JointDistributionNamed(
              {
                  # The autoregressive coefficients and the `log_scale` each follow
                  # an independent slow-moving random walk.
                  "coefs": independent.Independent(
                      normal.Normal(
                          loc=previous_state["coefs"], scale=0.01
                      ),
                      reinterpreted_batch_ndims=1,
                  ),
                  "log_scale": normal.Normal(
                      loc=previous_state["log_scale"], scale=0.01
                  ),
                  # The level is a linear combination of the previous *two* levels,
                  # with additional noise of scale `exp(log_scale)`.
                  "level": lambda coefs, log_scale: normal.Normal(  # pylint: disable=g-long-lambda
                      loc=(
                          coefs[..., 0] * previous_state["level"]
                          + coefs[..., 1]
                          * previous_state["previous_level"]
                      ),
                      scale=tf.exp(log_scale),
                  ),
                  # Store the previous level to access at the next step.
                  "previous_level": deterministic.Deterministic(
                      previous_state["level"]
                  ),
              }
          )
  
      process = markov_chain_lib.MarkovChain(
          # For simplicity, define the prior as a 'transition' from fixed values.
          initial_state_prior=transition_fn(
              0,
              previous_state={
                  "coefs": [0.7, -0.2],
                  "log_scale": -1.0,
                  "level": 0.0,
                  "previous_level": 0.0,
              },
          ),
          transition_fn=transition_fn,
          num_steps=100,
      )
      self.assertAllEqualNested(
          process.event_shape,
          {
              "coefs": [100, 2],
              "log_scale": [100],
              "level": [100],
              "previous_level": [100],
          },
      )
      self.assertAllEqual(
          process.batch_shape,
          {
              "coefs": [],
              "log_scale": [],
              "level": [],
              "previous_level": [],
          },
      )
  
      x = process.sample(5, seed=test_util.test_seed())
      self.assertAllEqual(x["coefs"].shape, [5, 100, 2])
      self.assertAllEqual(x["log_scale"].shape, [5, 100])
      self.assertAllEqual(x["level"].shape, [5, 100])
      self.assertAllEqual(x["previous_level"].shape, [5, 100])
  
      lp = process.log_prob(x)
      self.assertAllEqual(lp.shape, [5])
  
      x2, lp2 = process.experimental_sample_and_log_prob(
          2, seed=test_util.test_seed()
      )
      self.assertAllClose(lp2, process.log_prob(x2))
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.bijectors import split
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.distributions import markov_chain
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.distributions import wishart
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 6 for calling the MarkovChainTest obj's underlying member method code with example usage: obj.test_non_autobatched_joint_distribution() and returns: jdn