class_docstr: Hidden Markov model distribution.
  
  The `HiddenMarkovModel` distribution implements a (batch of) discrete hidden
  Markov models where the initial states, transition probabilities
  and observed states are all given by user-provided distributions.
  
  In this model, there is a sequence of integer-valued hidden states:
  `z[0], z[1], ..., z[num_steps - 1]` and a sequence of observed states:
  `x[0], ..., x[num_steps - 1]`.
  The distribution of `z[0]` is given by `initial_distribution`.
  The conditional probability of `z[i  +  1]` given `z[i]` is described by
  the batch of distributions in `transition_distribution`.
  For a batch of hidden Markov models, the coordinates before the rightmost one
  of the `transition_distribution` batch correspond to indices into the hidden
  Markov model batch. The rightmost coordinate of the batch is used to select
  which distribution `z[i + 1]` is drawn from.  The distributions corresponding
  to the probability of `z[i + 1]` conditional on `z[i] == k` is given by the
  elements of the batch whose rightmost coordinate is `k`.
  Similarly, the conditional distribution of `x[i]` given `z[i]` is given by
  the batch of `observation_distribution`.
  When the rightmost coordinate of `observation_distribution` is `k` it
  gives the conditional probabilities of `x[i]` given `z[i] == k`.
  The probability distribution associated with the `HiddenMarkovModel`
  distribution is the marginal distribution of `x[0],...,x[num_steps - 1]`.
  
  #### Examples
  
  ```python
  tfd = tfp.distributions
  
  # A simple weather model.
  
  # Represent a cold day with 0 and a hot day with 1.
  # Suppose the first day of a sequence has a 0.8 chance of being cold.
  # We can model this using the categorical distribution:
  
  initial_distribution = tfd.Categorical(probs=[0.8, 0.2])
  
  # Suppose a cold day has a 30% chance of being followed by a hot day
  # and a hot day has a 20% chance of being followed by a cold day.
  # We can model this as:
  
  transition_distribution = tfd.Categorical(probs=[[0.7, 0.3],
                                                   [0.2, 0.8]])
  
  # Suppose additionally that on each day the temperature is
  # normally distributed with mean and standard deviation 0 and 5 on
  # a cold day and mean and standard deviation 15 and 10 on a hot day.
  # We can model this with:
  
  observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])
  
  # We can combine these distributions into a single week long
  # hidden Markov model with:
  
  model = tfd.HiddenMarkovModel(
      initial_distribution=initial_distribution,
      transition_distribution=transition_distribution,
      observation_distribution=observation_distribution,
      num_steps=7)
  
  # The expected temperatures for each day are given by:
  
  model.mean()  # shape [7], elements approach 9.0
  
  # The log pdf of a week of temperature 0 is:
  
  model.log_prob(tf.zeros(shape=[7]))
  ```
  
  #### References
  [1] https://en.wikipedia.org/wiki/Hidden_Markov_model
class_docstr_tok: ["Hidden","Markov","model","distribution",".","The","`","HiddenMarkovModel","`","distribution","implements","a","(","batch","of",")","discrete","hidden","Markov","models","where","the","initial","states",",","transition","probabilities","and","observed","states","are","all","given","by","user-provided","distributions",".","In","this","model",",","there","is","a","sequence","of","integer-valued","hidden","states",":","`","z","[","0","]",",","z","[","1","]",",","...",",","z","[","num_steps","-","1","]","`","and","a","sequence","of","observed","states",":","`","x","[","0","]",",","...",",","x","[","num_steps","-","1","]","`",".","The","distribution","of","`","z","[","0","]","`","is","given","by","`","initial_distribution","`",".","The","conditional","probability","of","`","z","[","i","+","1","]","`","given","`","z","[","i","]","`","is","described","by","the","batch","of","distributions","in","`","transition_distribution","`",".","For","a","batch","of","hidden","Markov","models",",","the","coordinates","before","the","rightmost","one","of","the","`","transition_distribution","`","batch","correspond","to","indices","into","the","hidden","Markov","model","batch",".","The","rightmost","coordinate","of","the","batch","is","used","to","select","which","distribution","`","z","[","i","+","1","]","`","is","drawn","from",".","The","distributions","corresponding","to","the","probability","of","`","z","[","i","+","1","]","`","conditional","on","`","z","[","i","]","==","k","`","is","given","by","the","elements","of","the","batch","whose","rightmost","coordinate","is","`","k","`",".","Similarly",",","the","conditional","distribution","of","`","x","[","i","]","`","given","`","z","[","i","]","`","is","given","by","the","batch","of","`","observation_distribution","`",".","When","the","rightmost","coordinate","of","`","observation_distribution","`","is","`","k","`","it","gives","the","conditional","probabilities","of","`","x","[","i","]","`","given","`","z","[","i","]","==","k","`",".","The","probability","distribution","associated","with","the","`","HiddenMarkovModel","`","distribution","is","the","marginal","distribution","of","`","x","[","0","]",",","...",",","x","[","num_steps","-","1","]","`",".","#","#","#","#","Examples","``","`","python","tfd","=","tfp.distributions","#","A","simple","weather","model",".","#","Represent","a","cold","day","with","0","and","a","hot","day","with","1",".","#","Suppose","the","first","day","of","a","sequence","has","a","0.8","chance","of","being","cold",".","#","We","can","model","this","using","the","categorical","distribution",":","initial_distribution","=","tfd.Categorical","(","probs=","[","0.8",",","0.2","]",")","#","Suppose","a","cold","day","has","a","30","%","chance","of","being","followed","by","a","hot","day","#","and","a","hot","day","has","a","20","%","chance","of","being","followed","by","a","cold","day",".","#","We","can","model","this","as",":","transition_distribution","=","tfd.Categorical","(","probs=","[","[","0.7",",","0.3","]",",","[","0.2",",","0.8","]","]",")","#","Suppose","additionally","that","on","each","day","the","temperature","is","#","normally","distributed","with","mean","and","standard","deviation","0","and","5","on","#","a","cold","day","and","mean","and","standard","deviation","15","and","10","on","a","hot","day",".","#","We","can","model","this","with",":","observation_distribution","=","tfd.Normal","(","loc=","[","0.",",","15",".","]",",","scale=","[","5.",",","10",".","]",")","#","We","can","combine","these","distributions","into","a","single","week","long","#","hidden","Markov","model","with",":","model","=","tfd.HiddenMarkovModel","(","initial_distribution=initial_distribution",",","transition_distribution=transition_distribution",",","observation_distribution=observation_distribution",",","num_steps=7",")","#","The","expected","temperatures","for","each","day","are","given","by",":","model.mean","(",")","#","shape","[","7","]",",","elements","approach","9.0","#","The","log","pdf","of","a","week","of","temperature","0","is",":","model.log_prob","(","tf.zeros","(","shape=","[","7","]",")",")","``","`","#","#","#","#","References","[","1","]","https",":","\/\/en.wikipedia.org\/wiki\/Hidden_Markov_model"]
docstr: Computes forward pass.
  
  When observations `x[:T]` are given, returns the unnormalized filtering
    distribution `C(t) * p(z[t]| x[:t])` where `C(t)=p(x[:t])`.
  When observations are not given, returns prior probabilities `p(z[t]).`
  
  Args:
    observations: optional Tensor of observations.
    mask: optional Tensor of mask.
    log_transition: optional Tensor of log probabilities
      from the transition matrix with shape
      `[num_steps] + batch_shape + [num_states, num_states]`.
      If not specified, will be recomputed from self.transition_distribution.
  
  Returns:
    forward_log_probs: a Tensor of log probabilities with shape
      `[num_steps] + batch_shape + [num_states]`.
    observation_log_probs: a Tensor containing `log p(x[t])` with shape
      `[num_steps] + batch_shape`.
docstr_tok: ["Computes","forward","pass",".","When","observations","`","x","[",":","T","]","`","are","given",",","returns","the","unnormalized","filtering","distribution","`","C","(","t",")","*","p","(","z","[","t","]","|","x","[",":","t","]",")","`","where","`","C","(","t",")","=p","(","x","[",":","t","]",")","`",".","When","observations","are","not","given",",","returns","prior","probabilities","`","p","(","z","[","t","]",")",".","`","Args",":","observations",":","optional","Tensor","of","observations",".","mask",":","optional","Tensor","of","mask",".","log_transition",":","optional","Tensor","of","log","probabilities","from","the","transition","matrix","with","shape","`","[","num_steps","]","+","batch_shape","+","[","num_states",",","num_states","]","`",".","If","not","specified",",","will","be","recomputed","from","self.transition_distribution",".","Returns",":","forward_log_probs",":","a","Tensor","of","log","probabilities","with","shape","`","[","num_steps","]","+","batch_shape","+","[","num_states","]","`",".","observation_log_probs",":","a","Tensor","containing","`","log","p","(","x","[","t","]",")","`","with","shape","`","[","num_steps","]","+","batch_shape","`","."]
code: def _forward_pass(
      self, observations=None, mask=None, log_transition=None
  ):
      """Computes forward pass.
  
      When observations `x[:T]` are given, returns the unnormalized filtering
        distribution `C(t) * p(z[t]| x[:t])` where `C(t)=p(x[:t])`.
      When observations are not given, returns prior probabilities `p(z[t]).`
  
      Args:
        observations: optional Tensor of observations.
        mask: optional Tensor of mask.
        log_transition: optional Tensor of log probabilities
          from the transition matrix with shape
          `[num_steps] + batch_shape + [num_states, num_states]`.
          If not specified, will be recomputed from self.transition_distribution.
  
      Returns:
        forward_log_probs: a Tensor of log probabilities with shape
          `[num_steps] + batch_shape + [num_states]`.
        observation_log_probs: a Tensor containing `log p(x[t])` with shape
          `[num_steps] + batch_shape`.
      """
      mask = self._get_mask(mask)
      num_states = self.transition_distribution.batch_shape_tensor()[-1]
      log_init = _extract_log_probs(
          num_states,
          self.initial_distribution,
          move_time_dimension=False,
      )
      if log_transition is None:
          log_transition = _extract_log_probs(
              num_states,
              self.transition_distribution,
              move_time_dimension=self._time_varying_transition_distribution,
          )
  
      if observations is None:
          observation_log_probs = tf.zeros(
              ps.concat(
                  [
                      [self.num_steps],
                      self.batch_shape_tensor(),
                      [num_states],
                  ],
                  axis=0,
              ),
              dtype=log_init.dtype,
          )
      else:
          observation_tensor_shape = ps.shape(observations)
          observation_distribution = self.observation_distribution
          underlying_event_rank = ps.size(
              observation_distribution.event_shape_tensor()
          )
          mask_tensor_shape = (
              ps.shape(mask) if mask is not None else None
          )
          with self._observation_mask_shape_preconditions(
              observation_tensor_shape,
              mask_tensor_shape,
              underlying_event_rank,
          ):
              observation_log_probs = self._observation_log_probs(
                  observations, mask
              )
  
      log_prob = log_init + observation_log_probs[0]
  
      def _scan_multiple_steps_forwards():
          if self._time_varying_transition_distribution:
  
              def dynamic_forward_step(
                  log_prev_step, log_transition_and_observation
              ):
                  (
                      log_transition,
                      log_prob_observation,
                  ) = log_transition_and_observation
                  result = (
                      _log_vector_matrix(log_prev_step, log_transition)
                      + log_prob_observation
                  )
                  tensorshape_util.set_shape(
                      result, log_prob_observation.shape
                  )
                  return result
  
              forward_log_probs = tf.scan(
                  dynamic_forward_step,
                  (log_transition, observation_log_probs[1:]),
                  initializer=log_prob,
              )
          else:
  
              def forward_step(log_previous_step, log_prob_observation):
                  result = (
                      _log_vector_matrix(
                          log_previous_step, log_transition
                      )
                      + log_prob_observation
                  )
                  tensorshape_util.set_shape(
                      result, log_prob_observation.shape
                  )
                  return result
  
              forward_log_probs = tf.scan(
                  forward_step,
                  observation_log_probs[1:],
                  initializer=log_prob,
                  name="forward_log_probs",
              )
          return ps.concat([[log_prob], forward_log_probs], axis=0)
  
      def _scan_one_step_forwards():
          multiples = ps.concat(
              [[self.num_steps], ps.ones(ps.rank(log_prob), tf.int32)],
              axis=0,
          )
          return tf.tile(log_prob[tf.newaxis], multiples)
  
      forward_log_probs = ps.cond(
          self._num_steps > 1,
          _scan_multiple_steps_forwards,
          _scan_one_step_forwards,
      )
      return forward_log_probs, observation_log_probs
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import distribution
  from tensorflow_probability.python.internal import assert_util
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import reparameterization
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 22 for calling the HiddenMarkovModel obj's underlying member method code with example usage: obj._forward_pass(observations, mask, log_transition) and returns: forward_log_probs,  observation_log_probs, ps, tf, result, result