class_docstr: Applies distinct bijectors to initial + transition states of a chain.
docstr: Initializes the MarkovChain bijector.
  
  This bijector maps into the support of the corresponding MarkovChain
  distribution, using separate bijectors for the head (initial state)
  and tail (transition states) of the chain. Its input is a pair of
  unconstrained structures each matching `chain.dtype`, and the output is a
  constrained structure matching `chain.dtype`. Note that the inputs
  to the two bijectors may have different shapes, corresponding to the
  inverse images of the two bijectors, but the outputs must have the same
  shape in order to support concatenation along the `num_steps` axis.
  
  Conceptually, the Markov chain bijector performs the same transformation as
  the `joint_distribution._DefaultJointBijector` for a hypothetical joint
  distribution that samples from the chain one step at a time:
  
  ```python
  @tfd.JointDistributionCoroutineAutoBatched
  def markov_chain_equivalent():
    x = yield initial_state_prior
    for i in range(1, num_steps):
      x = yield transition_fn(i, x)
  markov_chain_equivalent_joint_bijector = (
    markov_chain_equivalent.experimental_default_event_space_bijector())
  ```
  
  However, just as `MarkovChain` uses low-level looping and batch operations
  for better performance than the corresponding joint
  distribution, the `MarkovChainBijector` is more efficient than the
  corresponding joint bijector.
  
  Args:
    chain: Instance of `tfd.MarkovChain`.
    bijector_fn: Callable with signature `bij = bijector_fn(dist)`, where
      `dist` is a `tfd.Distribution` instance. This is applied to the
      `chain.initial_state_prior` and to distributions returned by
      `chain.transition_fn(...)`.
    transition_bijector: Bijector instance for a single step of the
      transition model. This is typically equal to
      `bijector_fn(markov_chain.transition_fn(0,
      markov_chain.initial_state_prior.sample()))`; passing it explicitly
      avoids the need to recreate it whenever the chain bijector is
      copied or otherwise re-initialized.
    name: The name to give ops created by this bijector.
  
  #### Example
  
  For example, consider the following chain, which has dtype
  `{'probs': tf.float32}`, and describes a process in which a 2D vector
  is sampled from the probability simplex and then gradually corrupted by
  Gaussian noise (which will in general push it out of the simplex):
  
  ```python
  chain = tfd.MarkovChain(
    initial_state_prior=tfd.JointDistributionNamedAutoBatched(
        {'probs': tfd.Dirichlet([1., 1.])}),
    transition_fn=lambda _, x: tfd.JointDistributionNamedAutoBatched(
        {'probs': tfd.MultivariateNormalDiag(loc=x['probs'],
                                             scale_diag=[0.1, 0.1])},
      batch_ndims=ps.rank(x['probs'])),
    num_steps=10)
  ```
  
  Transformations of this distribution apply separate bijectors
  for the `Dirichlet` initial state and `MultivariateNormalDiag` transitions:
  
  ```python
  bij = chain.experimental_default_event_space_bijector()
  y = chain.sample(5)  # Shape: {'probs': [5, 10, 2]}
  x = chain.inverse(y)  # Shape: [{'probs': [5, 1]}, {'probs': [5, 9, 2]}]
  ```
  
  Note that the pulled-back `x` is a pair of structures, of shapes
  `{'probs': [5, 1]}` and  `{'probs': [5, 9, 2]}` respectively. The first is
  the result of pulling the initial state vectors back through the Dirichlet
  event space bijector, which inverts shape-`[2]` vectors on the simplex to
  shape-`[1]` unconstrained vectors.  The second comes from pulling the
  shape-`[2]` chain state at each of the remaining `9` timesteps back through
  the `MultivariateNormalDiag` event space bijector, which is just
  the identity bijector, resulting in vectors of shape `[2]`.
docstr_tok: ["Initializes","the","MarkovChain","bijector",".","This","bijector","maps","into","the","support","of","the","corresponding","MarkovChain","distribution",",","using","separate","bijectors","for","the","head","(","initial","state",")","and","tail","(","transition","states",")","of","the","chain",".","Its","input","is","a","pair","of","unconstrained","structures","each","matching","`","chain.dtype","`",",","and","the","output","is","a","constrained","structure","matching","`","chain.dtype","`",".","Note","that","the","inputs","to","the","two","bijectors","may","have","different","shapes",",","corresponding","to","the","inverse","images","of","the","two","bijectors",",","but","the","outputs","must","have","the","same","shape","in","order","to","support","concatenation","along","the","`","num_steps","`","axis",".","Conceptually",",","the","Markov","chain","bijector","performs","the","same","transformation","as","the","`","joint_distribution._DefaultJointBijector","`","for","a","hypothetical","joint","distribution","that","samples","from","the","chain","one","step","at","a","time",":","``","`","python","@","tfd.JointDistributionCoroutineAutoBatched","def","markov_chain_equivalent","(",")",":","x","=","yield","initial_state_prior","for","i","in","range","(","1",",","num_steps",")",":","x","=","yield","transition_fn","(","i",",","x",")","markov_chain_equivalent_joint_bijector","=","(","markov_chain_equivalent.experimental_default_event_space_bijector","(",")",")","``","`","However",",","just","as","`","MarkovChain","`","uses","low-level","looping","and","batch","operations","for","better","performance","than","the","corresponding","joint","distribution",",","the","`","MarkovChainBijector","`","is","more","efficient","than","the","corresponding","joint","bijector",".","Args",":","chain",":","Instance","of","`","tfd.MarkovChain","`",".","bijector_fn",":","Callable","with","signature","`","bij","=","bijector_fn","(","dist",")","`",",","where","`","dist","`","is","a","`","tfd.Distribution","`","instance",".","This","is","applied","to","the","`","chain.initial_state_prior","`","and","to","distributions","returned","by","`","chain.transition_fn","(","...",")","`",".","transition_bijector",":","Bijector","instance","for","a","single","step","of","the","transition","model",".","This","is","typically","equal","to","`","bijector_fn","(","markov_chain.transition_fn","(","0",",","markov_chain.initial_state_prior.sample","(",")",")",")","`",";","passing","it","explicitly","avoids","the","need","to","recreate","it","whenever","the","chain","bijector","is","copied","or","otherwise","re-initialized",".","name",":","The","name","to","give","ops","created","by","this","bijector",".","#","#","#","#","Example","For","example",",","consider","the","following","chain",",","which","has","dtype","`","{","'probs","'",":","tf.float32","}","`",",","and","describes","a","process","in","which","a","2D","vector","is","sampled","from","the","probability","simplex","and","then","gradually","corrupted","by","Gaussian","noise","(","which","will","in","general","push","it","out","of","the","simplex",")",":","``","`","python","chain","=","tfd.MarkovChain","(","initial_state_prior=tfd.JointDistributionNamedAutoBatched","(","{","'probs","'",":","tfd.Dirichlet","(","[","1.",",","1",".","]",")","}",")",",","transition_fn=lambda","_",",","x",":","tfd.JointDistributionNamedAutoBatched","(","{","'probs","'",":","tfd.MultivariateNormalDiag","(","loc=x","[","'probs","'","]",",","scale_diag=","[","0.1",",","0.1","]",")","}",",","batch_ndims=ps.rank","(","x","[","'probs","'","]",")",")",",","num_steps=10",")","``","`","Transformations","of","this","distribution","apply","separate","bijectors","for","the","`","Dirichlet","`","initial","state","and","`","MultivariateNormalDiag","`","transitions",":","``","`","python","bij","=","chain.experimental_default_event_space_bijector","(",")","y","=","chain.sample","(","5",")","#","Shape",":","{","'probs","'",":","[","5",",","10",",","2","]","}","x","=","chain.inverse","(","y",")","#","Shape",":","[","{","'probs","'",":","[","5",",","1","]","}",",","{","'probs","'",":","[","5",",","9",",","2","]","}","]","``","`","Note","that","the","pulled-back","`","x","`","is","a","pair","of","structures",",","of","shapes","`","{","'probs","'",":","[","5",",","1","]","}","`","and","`","{","'probs","'",":","[","5",",","9",",","2","]","}","`","respectively",".","The","first","is","the","result","of","pulling","the","initial","state","vectors","back","through","the","Dirichlet","event","space","bijector",",","which","inverts","shape-","`","[","2","]","`","vectors","on","the","simplex","to","shape-","`","[","1","]","`","unconstrained","vectors",".","The","second","comes","from","pulling","the","shape-","`","[","2","]","`","chain","state","at","each","of","the","remaining","`","9","`","timesteps","back","through","the","`","MultivariateNormalDiag","`","event","space","bijector",",","which","is","just","the","identity","bijector",",","resulting","in","vectors","of","shape","`","[","2","]","`","."]
code: def __init__(
      self,
      chain,
      bijector_fn,
      transition_bijector,
      name="markov_chain_bijector",
  ):
      """Initializes the MarkovChain bijector.
  
      This bijector maps into the support of the corresponding MarkovChain
      distribution, using separate bijectors for the head (initial state)
      and tail (transition states) of the chain. Its input is a pair of
      unconstrained structures each matching `chain.dtype`, and the output is a
      constrained structure matching `chain.dtype`. Note that the inputs
      to the two bijectors may have different shapes, corresponding to the
      inverse images of the two bijectors, but the outputs must have the same
      shape in order to support concatenation along the `num_steps` axis.
  
      Conceptually, the Markov chain bijector performs the same transformation as
      the `joint_distribution._DefaultJointBijector` for a hypothetical joint
      distribution that samples from the chain one step at a time:
  
      ```python
      @tfd.JointDistributionCoroutineAutoBatched
      def markov_chain_equivalent():
        x = yield initial_state_prior
        for i in range(1, num_steps):
          x = yield transition_fn(i, x)
      markov_chain_equivalent_joint_bijector = (
        markov_chain_equivalent.experimental_default_event_space_bijector())
      ```
  
      However, just as `MarkovChain` uses low-level looping and batch operations
      for better performance than the corresponding joint
      distribution, the `MarkovChainBijector` is more efficient than the
      corresponding joint bijector.
  
      Args:
        chain: Instance of `tfd.MarkovChain`.
        bijector_fn: Callable with signature `bij = bijector_fn(dist)`, where
          `dist` is a `tfd.Distribution` instance. This is applied to the
          `chain.initial_state_prior` and to distributions returned by
          `chain.transition_fn(...)`.
        transition_bijector: Bijector instance for a single step of the
          transition model. This is typically equal to
          `bijector_fn(markov_chain.transition_fn(0,
          markov_chain.initial_state_prior.sample()))`; passing it explicitly
          avoids the need to recreate it whenever the chain bijector is
          copied or otherwise re-initialized.
        name: The name to give ops created by this bijector.
  
      #### Example
  
      For example, consider the following chain, which has dtype
      `{'probs': tf.float32}`, and describes a process in which a 2D vector
      is sampled from the probability simplex and then gradually corrupted by
      Gaussian noise (which will in general push it out of the simplex):
  
      ```python
      chain = tfd.MarkovChain(
        initial_state_prior=tfd.JointDistributionNamedAutoBatched(
            {'probs': tfd.Dirichlet([1., 1.])}),
        transition_fn=lambda _, x: tfd.JointDistributionNamedAutoBatched(
            {'probs': tfd.MultivariateNormalDiag(loc=x['probs'],
                                                 scale_diag=[0.1, 0.1])},
          batch_ndims=ps.rank(x['probs'])),
        num_steps=10)
      ```
  
      Transformations of this distribution apply separate bijectors
      for the `Dirichlet` initial state and `MultivariateNormalDiag` transitions:
  
      ```python
      bij = chain.experimental_default_event_space_bijector()
      y = chain.sample(5)  # Shape: {'probs': [5, 10, 2]}
      x = chain.inverse(y)  # Shape: [{'probs': [5, 1]}, {'probs': [5, 9, 2]}]
      ```
  
      Note that the pulled-back `x` is a pair of structures, of shapes
      `{'probs': [5, 1]}` and  `{'probs': [5, 9, 2]}` respectively. The first is
      the result of pulling the initial state vectors back through the Dirichlet
      event space bijector, which inverts shape-`[2]` vectors on the simplex to
      shape-`[1]` unconstrained vectors.  The second comes from pulling the
      shape-`[2]` chain state at each of the remaining `9` timesteps back through
      the `MultivariateNormalDiag` event space bijector, which is just
      the identity bijector, resulting in vectors of shape `[2]`.
  
      """
      parameters = dict(locals())
      with tf.name_scope(name):
          self._chain = chain
          self._bijector_fn = bijector_fn
          self._initial_bijector = bijector_fn(
              chain.initial_state_prior
          )
          self._transition_bijector = transition_bijector
  
          inverse_min_event_ndims = tf.nest.map_structure(
              ps.rank_from_shape, chain.event_shape_tensor()
          )
          super(_MarkovChainBijector, self).__init__(
              forward_min_event_ndims=(
                  self._initial_bijector.inverse_event_ndims(
                      tf.nest.map_structure(
                          lambda nd: nd - 1, inverse_min_event_ndims
                      )
                  ),
                  self._transition_bijector.inverse_event_ndims(
                      inverse_min_event_ndims
                  ),
              ),
              inverse_min_event_ndims=inverse_min_event_ndims,
              is_constant_jacobian=(
                  self.initial_bijector.is_constant_jacobian
                  and self.transition_bijector.is_constant_jacobian
              ),
              validate_args=chain.validate_args,
              parameters=parameters,
              name=name,
          )
  
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 1 to create a new _MarkovChainBijector object from inherited base classes: bijector with example: obj = _MarkovChainBijector(chain, bijector_fn, transition_bijector, name)