class_docstr: Applies distinct bijectors to initial + transition states of a chain.
docstr: Runs the chain forward, accumulating `fn(b, x, y)` vals at every step.
  
  Args:
    fn: Callable with signature `result = fn(b, x, y)`.
    x0: Structure of initial state `Tensors`, each of shape
        `concat([[batch_shape], unconstrained_prior_event_shape])`.
    xs: Structure of `Tensors`, each of shape
        `concat([[batch_shape], [num_steps - 1],
        unconstrained_transition_event_shape])`.
  Returns:
    fs: Result `Tensor` of shape
      `concat([[num_steps], batch_shape, result_shape])`, where `result_shape`
      is the shape of the result from an unbatched call to `fn`.
docstr_tok: ["Runs","the","chain","forward",",","accumulating","`","fn","(","b",",","x",",","y",")","`","vals","at","every","step",".","Args",":","fn",":","Callable","with","signature","`","result","=","fn","(","b",",","x",",","y",")","`",".","x0",":","Structure","of","initial","state","`","Tensors","`",",","each","of","shape","`","concat","(","[","[","batch_shape","]",",","unconstrained_prior_event_shape","]",")","`",".","xs",":","Structure","of","`","Tensors","`",",","each","of","shape","`","concat","(","[","[","batch_shape","]",",","[","num_steps","-","1","]",",","unconstrained_transition_event_shape","]",")","`",".","Returns",":","fs",":","Result","`","Tensor","`","of","shape","`","concat","(","[","[","num_steps","]",",","batch_shape",",","result_shape","]",")","`",",","where","`","result_shape","`","is","the","shape","of","the","result","from","an","unbatched","call","to","`","fn","`","."]
code: def _apply_forward_scan(self, fn, x0, xs):
      """Runs the chain forward, accumulating `fn(b, x, y)` vals at every step.
  
      Args:
        fn: Callable with signature `result = fn(b, x, y)`.
        x0: Structure of initial state `Tensors`, each of shape
            `concat([[batch_shape], unconstrained_prior_event_shape])`.
        xs: Structure of `Tensors`, each of shape
            `concat([[batch_shape], [num_steps - 1],
            unconstrained_transition_event_shape])`.
      Returns:
        fs: Result `Tensor` of shape
          `concat([[num_steps], batch_shape, result_shape])`, where `result_shape`
          is the shape of the result from an unbatched call to `fn`.
      """
      xs_step_axes = tf.nest.map_structure(
          lambda nd: -nd,
          self.transition_bijector.inverse_event_ndims(
              # Outputs `y` have the num_steps axis at `-inverse_min_event_ndims`.
              self.inverse_min_event_ndims
          ),
      )
      xs = move_dimensions(xs, source=xs_step_axes, dest=0)
  
      # Evaluate the initial state.
      y0 = self.initial_bijector.forward(x0)
      f0 = fn(self.initial_bijector, x0, y0)
  
      # Evaluate the rest of the chain.
      def loop_body(previous_y_and_result, idx):
          previous_y, _ = previous_y_and_result
          bij = self.bijector_fn(
              self.chain.transition_fn(idx, previous_y)
          )
          x_i = tf.nest.map_structure(lambda x: x[idx - 1], xs)
          y_i = bij.forward(x_i)
          f_i = fn(bij, x_i, y_i)
          return (
              y_i,
              tf.nest.map_structure(
                  lambda a, b: tf.cast(a, b.dtype), f_i, f0
              ),
          )
  
      _, fs = tf.scan(
          loop_body,
          elems=tf.range(1, self.chain.num_steps),
          initializer=(y0, f0),
      )
      return concat_initial(f0, fs)
  
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 7 for calling the _MarkovChainBijector obj's underlying member method code with example usage: obj._apply_forward_scan(fn, x0, xs) and returns: concat_initial, y_i,  tf