class_docstr: Applies distinct bijectors to initial + transition states of a chain.
docstr: Applies `fn(b, y)` at each step of the chain.
  
  Args:
    fn: Callable with signature `result = fn(b, y)`.
    y: Structure of `Tensor`s, of shape
       `batch_shape + self.chain.event_shape`.
  Returns:
    f0: `Tensor` of shape `batch_shape + result_shape`.
    fs: `Tensor` of shape `[num_steps - 1] + batch_shape + result_shape`.
docstr_tok: ["Applies","`","fn","(","b",",","y",")","`","at","each","step","of","the","chain",".","Args",":","fn",":","Callable","with","signature","`","result","=","fn","(","b",",","y",")","`",".","y",":","Structure","of","`","Tensor","`","s",",","of","shape","`","batch_shape","+","self.chain.event_shape","`",".","Returns",":","f0",":","`","Tensor","`","of","shape","`","batch_shape","+","result_shape","`",".","fs",":","`","Tensor","`","of","shape","`","[","num_steps","-","1","]","+","batch_shape","+","result_shape","`","."]
code: def _apply_batch(self, fn, y):
      """Applies `fn(b, y)` at each step of the chain.
  
      Args:
        fn: Callable with signature `result = fn(b, y)`.
        y: Structure of `Tensor`s, of shape
           `batch_shape + self.chain.event_shape`.
      Returns:
        f0: `Tensor` of shape `batch_shape + result_shape`.
        fs: `Tensor` of shape `[num_steps - 1] + batch_shape + result_shape`.
      """
      y = move_dimensions(
          y, source=self.chain._step_axes(), dest=0
      )  # pylint: disable=protected-access
      f0 = fn(
          self.initial_bijector,
          tf.nest.map_structure(lambda y: y[0], y),
      )
      transition_dist = self.chain.transition_fn(
          tf.range(self.chain.num_steps - 1),
          tf.nest.map_structure(lambda y: y[:-1], y),
      )
      return (
          f0,
          fn(
              self.bijector_fn(transition_dist),
              tf.nest.map_structure(lambda y: y[1:], y),
          ),
      )
  
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 8 for calling the _MarkovChainBijector obj's underlying member method code with example usage: obj._apply_batch(fn, y) and returns: f0,  fn