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"]
code: def _sample_n(self, n, seed=None):
      init_seed, scan_seed, observation_seed = samplers.split_seed(
          seed, n=3, salt="HiddenMarkovModel"
      )
  
      transition_batch_shape = (
          self.transition_distribution.batch_shape_tensor()
      )
      num_states = transition_batch_shape[-1]
  
      batch_shape = self.batch_shape_tensor()
      batch_size = ps.reduce_prod(batch_shape)
      # The batch sizes of the underlying initial distributions and
      # transition distributions might not match the batch size of
      # the HMM distribution.
      # As a result we need to ask for more samples from the
      # underlying distributions and then reshape the results into
      # the correct batch size for the HMM.
      init_repeat = ps.reduce_prod(batch_shape) // ps.reduce_prod(
          self._initial_distribution.batch_shape_tensor()
      )
      init_state = self._initial_distribution.sample(
          n * init_repeat, seed=init_seed
      )
      init_state = tf.reshape(init_state, [n, batch_size])
      # init_state :: n batch_size
  
      transition_repeat = ps.reduce_prod(batch_shape) // ps.reduce_prod(
          transition_batch_shape[:-2]
          if self._time_varying_transition_distribution
          else transition_batch_shape[:-1]
      )
  
      init_shape = init_state.shape
  
      def generate_step(state_and_seed, step):
          """Take a single step in Markov chain."""
          state, seed = state_and_seed
          sample_seed, next_seed = samplers.split_seed(seed)
  
          if self._time_varying_transition_distribution:
              gen = self._transition_distribution[..., step, :].sample(
                  n * transition_repeat, seed=sample_seed
              )
          else:
              gen = self._transition_distribution.sample(
                  n * transition_repeat, seed=sample_seed
              )
          # gen :: (n * transition_repeat) transition_batch
  
          new_states = tf.reshape(gen, [n, batch_size, num_states])
  
          # new_states :: n batch_size num_states
  
          old_states_one_hot = tf.one_hot(
              state, num_states, dtype=tf.int32
          )
  
          # old_states :: n batch_size num_states
  
          result = tf.reduce_sum(
              old_states_one_hot * new_states, axis=-1
          )
          # We know that `generate_step` must preserve the shape of the
          # tensor of states of each state. This is because
          # the transition matrix must be square. But TensorFlow might
          # not know this so we explicitly tell it that the result has the
          # same shape.
          tensorshape_util.set_shape(result, init_shape)
          return result, next_seed
  
      def _scan_multiple_steps():
          """Take multiple steps with tf.scan."""
          step = ps.range(self._num_steps - 1, dtype=tf.int32)
          hidden_states, _ = tf.scan(
              generate_step, step, initializer=(init_state, scan_seed)
          )
  
          # TODO(b/115618503): add/use prepend_initializer to tf.scan
          return tf.concat([[init_state], hidden_states], axis=0)
  
      def _scan_one_step():
          multiples = ps.concat(
              [
                  [self.num_steps],
                  ps.ones(ps.rank(init_state), tf.int32),
              ],
              axis=0,
          )
          return tf.tile(init_state[tf.newaxis, ...], multiples)
  
      hidden_states = ps.cond(
          self._num_steps > 1, _scan_multiple_steps, _scan_one_step
      )
  
      hidden_one_hot = tf.one_hot(
          hidden_states,
          num_states,
          dtype=self._observation_distribution.dtype,
      )
      # hidden_one_hot :: num_steps n batch_size num_states
  
      # The observation distribution batch size might not match
      # the required batch size so as with the initial and
      # transition distributions we generate more samples and
      # reshape.
      observation_repeat = ps.maximum(
          batch_size
          // ps.reduce_prod(
              self._observation_distribution.batch_shape_tensor()[:-1]
          ),
          1,
      )
  
      if self._time_varying_observation_distribution:
          possible_observations = self._observation_distribution.sample(
              [observation_repeat * n], seed=observation_seed
          )
          # possible observations needs to have num_steps moved to the beginning.
          possible_observations = distribution_util.move_dimension(
              possible_observations,
              -(
                  tf.size(
                      self._observation_distribution.event_shape_tensor()
                  )
                  + 2
              ),
              0,
          )
      else:
          possible_observations = self._observation_distribution.sample(
              [self._num_steps, observation_repeat * n],
              seed=observation_seed,
          )
  
      inner_shape = self._observation_distribution.event_shape_tensor()
  
      # possible_observations :: num_steps (observation_repeat * n)
      #                          observation_batch[:-1] num_states inner_shape
  
      possible_observations = tf.reshape(
          possible_observations,
          ps.concat(
              [
                  [self._num_steps, n],
                  batch_shape,
                  [num_states],
                  inner_shape,
              ],
              axis=0,
          ),
      )
  
      # possible_observations :: steps n batch_size num_states inner_shape
  
      hidden_one_hot = tf.reshape(
          hidden_one_hot,
          ps.concat(
              [
                  [self._num_steps, n],
                  batch_shape,
                  [num_states],
                  ps.ones_like(inner_shape),
              ],
              axis=0,
          ),
      )
  
      # hidden_one_hot :: steps n batch_size num_states "inner_shape"
  
      observations = tf.reduce_sum(
          hidden_one_hot * possible_observations,
          axis=-1 - ps.size(inner_shape),
      )
      # observations :: steps n batch_size inner_shape
  
      observations = distribution_util.move_dimension(
          observations, 0, 1 + ps.size(batch_shape)
      )
      # returned :: n batch_shape steps inner_shape
  
      return observations
  
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 13 for calling the HiddenMarkovModel obj's underlying member method code with example usage: obj._sample_n(n, seed) and returns: observations, result,  next_seed, tf, tf