class_docstr: Distribution of a sequence generated by a memoryless process.
  
  A discrete-time [Markov chain](https://en.wikipedia.org/wiki/Markov_chain)
  is a sequence of random variables in which the variable(s) at each step is
  independent of all previous variables, *conditioned on* the variable(s) at the
  immediate predecessor step. That is, there can be no (direct) long-term
  dependencies. This 'Markov property' is a simplifying assumption; for example,
  it enables efficient sampling. Many time-series models can be formulated as
  Markov chains.
  
  Instances of `tfd.MarkovChain` represent fully-observed, discrete-time Markov
  chains, with one or more random variables at each step. These variables may
  take continuous or discrete values. Sampling is done sequentially, requiring
  time that scales with the length of the sequence; `log_prob` evaluation is
  vectorized over timesteps, and so requires only constant time given sufficient
  parallelism.
  
  #### Related distributions
  
  The discrete-valued Markov chains modeled by `tfd.HiddenMarkovModel` (using
  a trivial observation distribution) are a special case of those supported by
  this distribution, which enable exact inference over the values in an
  unobserved chain. Continuous-valued chains with linear Gaussian transitions
  are supported by `tfd.LinearGaussianStateSpaceModel`, which can similarly
  exploit the linear Gaussian structure for exact inference of hidden states.
  These distributions are limited to chains that have the respective (discrete
  or linear Gaussian) structure.
  
  Autoregressive models that do *not* necessarily respect the Markov property
  are supported by `tfd.Autoregressive`, which is, in that sense, more general
  than this distribution. These models require a more involved specification,
  and sampling in general requires quadratic (rather than linear) time in the
  length of the sequence.
  
  Exact inference for unobserved Markov chains is not possible in
  general; however, particle filtering exploits the Markov property
  to perform approximate inference, and is often a well-suited method for
  sequential inference tasks. Particle filtering is available in TFP using
  `tfp.experimental.mcmc.particle_filter`, and related methods.
  
  #### Example: Gaussian random walk
  
  One of the simplest continuous-valued Markov chains is a
  [Gaussian random walk](
  https://en.wikipedia.org/wiki/Random_walk#Gaussian_random_walk).
  This may also be viewed as a discretized [Brownian motion](
  https://en.wikipedia.org/wiki/Brownian_motion).
  
  ```python
  tfd = tfp.distributions
  
  gaussian_walk = tfd.MarkovChain(
    initial_state_prior=tfd.Normal(loc=0., scale=1.),
    transition_fn=lambda _, x: tfd.Normal(loc=x, scale=1.),
    num_steps=100)
  # ==> `gaussian_walk.event_shape == [100]`
  # ==> `gaussian_walk.batch_shape == []`
  
  x = gaussian_walk.sample(5)  # Samples a matrix of 5 independent walks.
  lp = gaussian_walk.log_prob(x)  # ==> `lp.shape == [5]`.
  ```
  
  #### Example: batch of random walks
  
  To spice things up, we'll now define a *batch* of random walks, each following
  a different distribution (in this case, different starting locations).
  We'll also demonstrate scales that differ across timesteps.
  
  ```python
  scales = tf.convert_to_tensor([0.5, 0.3, 0.2, 0.2, 0.3, 0.2, 0.7])
  batch_gaussian_walk = tfd.MarkovChain(
    # The prior distribution determines the batch shape for the chain.
    # Transitions must respect this batch shape.
    initial_state_prior=tfd.Normal(loc=[-10., 0., 10.],
                                   scale=[1., 1., 1.]),
    transition_fn=lambda t, x: tfd.Normal(
      loc=x,
      # The `num_steps` dimension will always be leftmost in `x`, so we
      # pad the scale to the same rank as `x` to make their shapes line up.
      tf.reshape(tf.gather(scales, t),
                 tf.concat([[-1],
                            tf.ones(tf.rank(x) - 1, dtype=tf.int32)], axis=0))),
    # Limit to eight steps since we only specified scales for seven transitions.
    num_steps=8)
  # ==> `batch_gaussian_walk.event_shape == [8]`
  # ==> `batch_gaussian_walk.batch_shape == [3]`
  
  x = batch_gaussian_walk.sample(5)  # ==> `x.shape == [5, 3, 8]`.
  lp = batch_gaussian_walk.log_prob(x)  # ==> `lp.shape == [5, 3]`.
  ```
  
  #### Example: multivariate chain with longer-term dependence
  
  We can also define multivariate Markov chains. In addition to the obvious
  use of modeling the joint evolution of multiple variables, multivariate
  chains can also help us work around the Markov limitation by
  the trick of folding state history into the current state as an auxiliary
  variable(s). The next example, a second-order [autoregressive process](
  https://en.wikipedia.org/wiki/Autoregressive_model) with dynamic coefficients
  and scale, contains multiple time-dependent variables and also uses an
  auxiliary `previous_level` variable to enable the transition function
  to access the previous *two* steps of history:
  
  ```python
  
  def transition_fn(_, previous_state):
    return tfd.JointDistributionNamedAutoBatched(
        # The transition distribution must match the batch shape of the chain.
        # Since `log_scale` is a scalar quantity, its shape is the batch shape.
        batch_ndims=tf.rank(previous_state['log_scale']),
        model={
            # The autoregressive coefficients and the `log_scale` each follow
            # an independent slow-moving random walk.
            'coefs': tfd.Normal(loc=previous_state['coefs'], scale=0.01),
            'log_scale': tfd.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: tfd.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': tfd.Deterministic(previous_state['level'])})
  ```
  
  Note: when using an autobatched joint distribution as a transition model,
  as we did here, it is necessary to explicitly set its `batch_ndims` to the
  batch rank of the passed-in state. This will be at least the batch rank of the
  initial state prior, but may be greater, e.g., when evaluating multiple iid
  samples. In general, the correct batch rank is that of the previous state
  `Tensor`s.
  
  ```python
  process = tfd.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.,
              'level': 0.,
              'previous_level': 0.}),
      transition_fn=transition_fn,
      num_steps=100)
  # ==> `process.event_shape == {'coefs': [100, 2], 'log_scale': [100],
  #                              'level': [100], 'previous_level': [100]}`
  # ==> `process.batch_shape == []`
  
  x = process.sample(5)
  # ==> `x['coefs'].shape == [5, 100, 2]`
  # ==> `x['log_scale'].shape == [5, 100]`
  # ==> `x['level'].shape == [5, 100]`
  # ==> `x['previous_level'].shape == [5, 100]`
  lp = process.log_prob(x)  # ==> `lp.shape == [5]`.
  ```
code: def _sample_control_dependencies(self, x):
      if not self.validate_args:
          return []
      parts_num_steps = tf.nest.flatten(
          tf.nest.map_structure(
              lambda x, k: ps.shape(x)[k], x, self._step_axes()
          )
      )
      return [
          assert_util.assert_equal(  # pylint: disable=g-complex-comprehension
              num_steps,
              self.num_steps,
              message="Input shape does not match the expected num_steps.",
          )
          for num_steps in parts_num_steps
      ]
  
imports: import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.distributions import log_prob_ratio
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.math import generic
label_desc: Use image node_id 19 for calling the MarkovChain obj's underlying member method code with example usage: obj._sample_control_dependencies(x) and returns: unknown, list