class_docstr: HMM with discrete latent states and discrete observations, allowing for
  missing data or variable length sequences. Observations are assumed
  to be one hot encoded; rows with all zeros indicate missing data.
  
  .. warning:: Unlike in pyro's pyro.distributions.DiscreteHMM, which
      computes the probability of the first state as
      initial.T @ transition @ emission
      this distribution uses the standard HMM convention,
      initial.T @ emission
  
  :param ~torch.Tensor initial_logits: A logits tensor for an initial
      categorical distribution over latent states. Should have rightmost
      size ``state_dim`` and be broadcastable to
      ``(batch_size, state_dim)``.
  :param ~torch.Tensor transition_logits: A logits tensor for transition
      conditional distributions between latent states. Should have rightmost
      shape ``(state_dim, state_dim)`` (old, new), and be broadcastable
      to ``(batch_size, state_dim, state_dim)``.
  :param ~torch.Tensor observation_logits: A logits tensor for observation
      distributions from latent states. Should have rightmost shape
      ``(state_dim, categorical_size)``, where ``categorical_size`` is the
      dimension of the categorical output, and be broadcastable
      to ``(batch_size, state_dim, categorical_size)``.
class_docstr_tok: ["HMM","with","discrete","latent","states","and","discrete","observations",",","allowing","for","missing","data","or","variable","length","sequences",".","Observations","are","assumed","to","be","one","hot","encoded",";","rows","with","all","zeros","indicate","missing","data",".","..","warning",":",":","Unlike","in","pyro","'s","pyro.distributions.DiscreteHMM",",","which","computes","the","probability","of","the","first","state","as","initial.T","@","transition","@","emission","this","distribution","uses","the","standard","HMM","convention",",","initial.T","@","emission",":","param","~torch.Tensor","initial_logits",":","A","logits","tensor","for","an","initial","categorical","distribution","over","latent","states",".","Should","have","rightmost","size","``","state_dim","``","and","be","broadcastable","to","``","(","batch_size",",","state_dim",")","``",".",":","param","~torch.Tensor","transition_logits",":","A","logits","tensor","for","transition","conditional","distributions","between","latent","states",".","Should","have","rightmost","shape","``","(","state_dim",",","state_dim",")","``","(","old",",","new",")",",","and","be","broadcastable","to","``","(","batch_size",",","state_dim",",","state_dim",")","``",".",":","param","~torch.Tensor","observation_logits",":","A","logits","tensor","for","observation","distributions","from","latent","states",".","Should","have","rightmost","shape","``","(","state_dim",",","categorical_size",")","``",",","where","``","categorical_size","``","is","the","dimension","of","the","categorical","output",",","and","be","broadcastable","to","``","(","batch_size",",","state_dim",",","categorical_size",")","``","."]
docstr: Distribution conditional on the state variable.
  
  :param ~torch.Tensor map_states: State trajectory. Must be
      integer-valued (long) and broadcastable to
      ``(batch_size, num_steps)``.
docstr_tok: ["Distribution","conditional","on","the","state","variable",".",":","param","~torch.Tensor","map_states",":","State","trajectory",".","Must","be","integer-valued","(","long",")","and","broadcastable","to","``","(","batch_size",",","num_steps",")","``","."]
code: def given_states(self, states):
      """
      Distribution conditional on the state variable.
  
      :param ~torch.Tensor map_states: State trajectory. Must be
          integer-valued (long) and broadcastable to
          ``(batch_size, num_steps)``.
      """
      shape = broadcast_shape(
          list(self.batch_shape) + [1, 1],
          list(states.shape[:-1]) + [1, 1],
          [1, 1, self.observation_logits.shape[-1]],
      )
      states_index = states.unsqueeze(-1) * torch.ones(
          shape, dtype=torch.long
      )
      obs_logits = self.observation_logits * torch.ones(shape)
      logits = torch.gather(obs_logits, -2, states_index)
      return OneHotCategorical(logits=logits)
  
file_path: pyro/pyro/contrib/mue/missingdatahmm.py
imports: import torch
  from torch.distributions import Categorical, OneHotCategorical
  from pyro.distributions import constraints
  from pyro.distributions.hmm import _sequential_logmatmulexp
  from pyro.distributions.torch_distribution import TorchDistribution
  from pyro.distributions.util import broadcast_shape
label_desc: Use image node_id 8 for calling the MissingDataDiscreteHMM obj's underlying member method code with example usage: obj.given_states(states) and returns: OneHotCategorical