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: Compute maximum likelihood sequence of hidden states.
  
  When this function is provided with a sequence of observations
  `x[0], ..., x[num_steps - 1]`, it returns the sequence of hidden
  states `z[0], ..., z[num_steps - 1]`, drawn from the underlying
  Markov chain, that is most likely to yield those observations.
  
  It uses the [Viterbi algorithm](
  https://en.wikipedia.org/wiki/Viterbi_algorithm).
  
  Note: the behavior of this function is undefined if the
  `observations` argument represents impossible observations
  from the model.
  
  Note: if there isn't a unique most likely sequence then one
  of the equally most likely sequences is chosen.
  
  Args:
    observations: A tensor representing a batch of observations made on the
      hidden Markov model.  The rightmost dimensions of this tensor correspond
      to the dimensions of the observation distributions of the underlying
      Markov chain, if the observations are non-scalar.  The next dimension
      from the right indexes the steps in a sequence of observations from a
      single sample from the hidden Markov model.  The size of this dimension
      should match the `num_steps` parameter of the hidden Markov model
      object.  The other dimensions are the dimensions of the batch and these
      are broadcast with the hidden Markov model's parameters.
    mask: optional bool-type `tensor` with rightmost dimension matching
      `num_steps`, indicating which observations should be ignored
      (not conditioned on). When the mask has value `True`, the corresponding
      observations aren't used. If no mask is specified (`mask` and
      `self.mask` are both `None`) then all of the observations are used. The
      leftmost dimensions `shape(mask)[:-1]` must broadcast with
      `self.batch_shape`.
    name: Python `str` name prefixed to Ops created by this class.
      Default value: "HiddenMarkovModel".
  
  Returns:
    posterior_mode: A `Tensor` representing the most likely sequence of hidden
      states. The rightmost dimension of this tensor will equal the
      `num_steps` parameter providing one hidden state for each step. The
      other dimensions are those of the batch.
  
  Raises:
    ValueError: if the `observations` tensor does not consist of
    sequences of `num_steps` observations.
  
  #### 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.
  
  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.
  
  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.
  
  observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])
  
  # This gives the hidden Markov model:
  
  model = tfd.HiddenMarkovModel(
      initial_distribution=initial_distribution,
      transition_distribution=transition_distribution,
      observation_distribution=observation_distribution,
      num_steps=7)
  
  # Suppose we observe gradually rising temperatures over a week:
  temps = [-2., 0., 2., 4., 6., 8., 10.]
  
  # We can now compute the most probable sequence of hidden states:
  
  model.posterior_mode(temps)
  
  # The result is [0 0 0 0 0 1 1] telling us that the transition
  # from "cold" to "hot" most likely happened between the
  # 5th and 6th days.
  ```
docstr_tok: ["Compute","maximum","likelihood","sequence","of","hidden","states",".","When","this","function","is","provided","with","a","sequence","of","observations","`","x","[","0","]",",","...",",","x","[","num_steps","-","1","]","`",",","it","returns","the","sequence","of","hidden","states","`","z","[","0","]",",","...",",","z","[","num_steps","-","1","]","`",",","drawn","from","the","underlying","Markov","chain",",","that","is","most","likely","to","yield","those","observations",".","It","uses","the","[","Viterbi","algorithm","]","(","https",":","\/\/en.wikipedia.org\/wiki\/Viterbi_algorithm",")",".","Note",":","the","behavior","of","this","function","is","undefined","if","the","`","observations","`","argument","represents","impossible","observations","from","the","model",".","Note",":","if","there","is","n't","a","unique","most","likely","sequence","then","one","of","the","equally","most","likely","sequences","is","chosen",".","Args",":","observations",":","A","tensor","representing","a","batch","of","observations","made","on","the","hidden","Markov","model",".","The","rightmost","dimensions","of","this","tensor","correspond","to","the","dimensions","of","the","observation","distributions","of","the","underlying","Markov","chain",",","if","the","observations","are","non-scalar",".","The","next","dimension","from","the","right","indexes","the","steps","in","a","sequence","of","observations","from","a","single","sample","from","the","hidden","Markov","model",".","The","size","of","this","dimension","should","match","the","`","num_steps","`","parameter","of","the","hidden","Markov","model","object",".","The","other","dimensions","are","the","dimensions","of","the","batch","and","these","are","broadcast","with","the","hidden","Markov","model","'s","parameters",".","mask",":","optional","bool-type","`","tensor","`","with","rightmost","dimension","matching","`","num_steps","`",",","indicating","which","observations","should","be","ignored","(","not","conditioned","on",")",".","When","the","mask","has","value","`","True","`",",","the","corresponding","observations","are","n't","used",".","If","no","mask","is","specified","(","`","mask","`","and","`","self.mask","`","are","both","`","None","`",")","then","all","of","the","observations","are","used",".","The","leftmost","dimensions","`","shape","(","mask",")","[",":","-1","]","`","must","broadcast","with","`","self.batch_shape","`",".","name",":","Python","`","str","`","name","prefixed","to","Ops","created","by","this","class",".","Default","value",":","``","HiddenMarkovModel","''",".","Returns",":","posterior_mode",":","A","`","Tensor","`","representing","the","most","likely","sequence","of","hidden","states",".","The","rightmost","dimension","of","this","tensor","will","equal","the","`","num_steps","`","parameter","providing","one","hidden","state","for","each","step",".","The","other","dimensions","are","those","of","the","batch",".","Raises",":","ValueError",":","if","the","`","observations","`","tensor","does","not","consist","of","sequences","of","`","num_steps","`","observations",".","#","#","#","#","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",".","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",".","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",".","observation_distribution","=","tfd.Normal","(","loc=","[","0.",",","15",".","]",",","scale=","[","5.",",","10",".","]",")","#","This","gives","the","hidden","Markov","model",":","model","=","tfd.HiddenMarkovModel","(","initial_distribution=initial_distribution",",","transition_distribution=transition_distribution",",","observation_distribution=observation_distribution",",","num_steps=7",")","#","Suppose","we","observe","gradually","rising","temperatures","over","a","week",":","temps","=","[","-2.",",","0.",",","2.",",","4.",",","6.",",","8.",",","10",".","]","#","We","can","now","compute","the","most","probable","sequence","of","hidden","states",":","model.posterior_mode","(","temps",")","#","The","result","is","[","0","0","0","0","0","1","1","]","telling","us","that","the","transition","#","from","``","cold","''","to","``","hot","''","most","likely","happened","between","the","#","5th","and","6th","days",".","``","`"]
code: def posterior_mode(
      self, observations, mask=None, name="posterior_mode"
  ):
      """Compute maximum likelihood sequence of hidden states.
  
      When this function is provided with a sequence of observations
      `x[0], ..., x[num_steps - 1]`, it returns the sequence of hidden
      states `z[0], ..., z[num_steps - 1]`, drawn from the underlying
      Markov chain, that is most likely to yield those observations.
  
      It uses the [Viterbi algorithm](
      https://en.wikipedia.org/wiki/Viterbi_algorithm).
  
      Note: the behavior of this function is undefined if the
      `observations` argument represents impossible observations
      from the model.
  
      Note: if there isn't a unique most likely sequence then one
      of the equally most likely sequences is chosen.
  
      Args:
        observations: A tensor representing a batch of observations made on the
          hidden Markov model.  The rightmost dimensions of this tensor correspond
          to the dimensions of the observation distributions of the underlying
          Markov chain, if the observations are non-scalar.  The next dimension
          from the right indexes the steps in a sequence of observations from a
          single sample from the hidden Markov model.  The size of this dimension
          should match the `num_steps` parameter of the hidden Markov model
          object.  The other dimensions are the dimensions of the batch and these
          are broadcast with the hidden Markov model's parameters.
        mask: optional bool-type `tensor` with rightmost dimension matching
          `num_steps`, indicating which observations should be ignored
          (not conditioned on). When the mask has value `True`, the corresponding
          observations aren't used. If no mask is specified (`mask` and
          `self.mask` are both `None`) then all of the observations are used. The
          leftmost dimensions `shape(mask)[:-1]` must broadcast with
          `self.batch_shape`.
        name: Python `str` name prefixed to Ops created by this class.
          Default value: "HiddenMarkovModel".
  
      Returns:
        posterior_mode: A `Tensor` representing the most likely sequence of hidden
          states. The rightmost dimension of this tensor will equal the
          `num_steps` parameter providing one hidden state for each step. The
          other dimensions are those of the batch.
  
      Raises:
        ValueError: if the `observations` tensor does not consist of
        sequences of `num_steps` observations.
  
      #### 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.
  
      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.
  
      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.
  
      observation_distribution = tfd.Normal(loc=[0., 15.], scale=[5., 10.])
  
      # This gives the hidden Markov model:
  
      model = tfd.HiddenMarkovModel(
          initial_distribution=initial_distribution,
          transition_distribution=transition_distribution,
          observation_distribution=observation_distribution,
          num_steps=7)
  
      # Suppose we observe gradually rising temperatures over a week:
      temps = [-2., 0., 2., 4., 6., 8., 10.]
  
      # We can now compute the most probable sequence of hidden states:
  
      model.posterior_mode(temps)
  
      # The result is [0 0 0 0 0 1 1] telling us that the transition
      # from "cold" to "hot" most likely happened between the
      # 5th and 6th days.
      ```
      """
  
      with self._name_and_control_scope(name):
          mask = self._get_mask(mask)
          mask_tensor_shape = (
              ps.shape(mask) if mask is not None else None
          )
          observations = tf.convert_to_tensor(
              observations, name="observations"
          )
          num_states = (
              self.transition_distribution.batch_shape_tensor()[-1]
          )
          observation_distribution = self.observation_distribution
          underlying_event_rank = ps.size(
              observation_distribution.event_shape_tensor()
          )
          observation_tensor_shape = ps.shape(observations)
          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_init = _extract_log_probs(
                  num_states,
                  self.initial_distribution,
                  move_time_dimension=False,
              )
              log_trans = _extract_log_probs(
                  num_states,
                  self.transition_distribution,
                  move_time_dimension=self._time_varying_transition_distribution,
              )
  
              log_prob = log_init + observation_log_probs[0]
  
              def _reduce_multiple_steps():
                  """Perform `reduce_max` operation when `num_steps` > 1."""
  
                  if self._time_varying_transition_distribution:
  
                      def dynamic_forward_step(
                          previous_step_pair, log_trans_and_obs
                      ):
                          (
                              log_trans,
                              log_prob_observation,
                          ) = log_trans_and_obs
                          log_prob_previous = previous_step_pair[0]
                          log_prob = (
                              log_prob_previous[..., tf.newaxis]
                              + log_trans
                              + log_prob_observation[..., tf.newaxis, :]
                          )
                          most_likely_given_successor = tf.argmax(
                              log_prob, axis=-2
                          )
                          max_log_p_given_successor = tf.reduce_max(
                              log_prob, axis=-2
                          )
                          return (
                              max_log_p_given_successor,
                              most_likely_given_successor,
                          )
  
                      (
                          forward_log_probs,
                          all_most_likely_given_successor,
                      ) = tf.scan(
                          dynamic_forward_step,
                          (log_trans, observation_log_probs[1:]),
                          initializer=(
                              log_prob,
                              tf.zeros(
                                  ps.shape(log_prob), dtype=tf.int64
                              ),
                          ),
                          name="forward_log_probs",
                      )
                  else:
  
                      def forward_step(
                          previous_step_pair, log_prob_observation
                      ):
                          log_prob_previous = previous_step_pair[0]
                          log_prob = (
                              log_prob_previous[..., tf.newaxis]
                              + log_trans
                              + log_prob_observation[..., tf.newaxis, :]
                          )
                          most_likely_given_successor = tf.argmax(
                              log_prob, axis=-2
                          )
                          max_log_p_given_successor = tf.reduce_max(
                              log_prob, axis=-2
                          )
                          return (
                              max_log_p_given_successor,
                              most_likely_given_successor,
                          )
  
                      (
                          forward_log_probs,
                          all_most_likely_given_successor,
                      ) = tf.scan(
                          forward_step,
                          observation_log_probs[1:],
                          initializer=(
                              log_prob,
                              tf.zeros(
                                  ps.shape(log_prob), dtype=tf.int64
                              ),
                          ),
                          name="forward_log_probs",
                      )
  
                  most_likely_end = tf.argmax(
                      forward_log_probs[-1], axis=-1
                  )
  
                  # We require the operation that gives C from A and B where
                  # C[i...j] = A[i...j, B[i...j]]
                  # and A = most_likely_given_successor
                  #     B = most_likely_successor.
                  # tf.gather requires indices of known shape so instead we use
                  # reduction with tf.one_hot(B) to pick out elements from B
                  def backward_step(
                      most_likely_successor, most_likely_given_successor
                  ):
                      return tf.reduce_sum(
                          (
                              most_likely_given_successor
                              * tf.one_hot(
                                  most_likely_successor,
                                  num_states,
                                  dtype=tf.int64,
                              )
                          ),
                          axis=-1,
                      )
  
                  backward_scan = tf.scan(
                      backward_step,
                      all_most_likely_given_successor,
                      most_likely_end,
                      reverse=True,
                  )
                  most_likely_sequences = tf.concat(
                      [backward_scan, [most_likely_end]], axis=0
                  )
                  return distribution_util.move_dimension(
                      most_likely_sequences, 0, -1
                  )
  
              def _reduce_one_step():
                  multiples = ps.concat(
                      [
                          ps.ones(ps.rank(log_prob) - 1, tf.int32),
                          [self.num_steps],
                      ],
                      axis=0,
                  )
                  return tf.tile(
                      tf.argmax(log_prob, axis=-1)[..., tf.newaxis],
                      multiples,
                  )
  
              return ps.cond(
                  self.num_steps > 1,
                  _reduce_multiple_steps,
                  _reduce_one_step,
              )
  
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 24 for calling the HiddenMarkovModel obj's underlying member method code with example usage: obj.posterior_mode(observations, mask, name) and returns: ps, distribution_util, tf, tf, max_log_p_given_successor,  most_likely_given_successor, max_log_p_given_successor,  most_likely_given_successor