class_docstr: Hidden Markov Model with discrete latent state and arbitrary observation
  distribution.
  
  This uses [1] to parallelize over time, achieving O(log(time)) parallel
  complexity for computing :meth:`log_prob`, :meth:`filter`, and :meth:`sample`.
  
  The event_shape of this distribution includes time on the left::
  
      event_shape = (num_steps,) + observation_dist.event_shape
  
  This distribution supports any combination of homogeneous/heterogeneous
  time dependency of ``transition_logits`` and ``observation_dist``. However,
  because time is included in this distribution's event_shape, the
  homogeneous+homogeneous case will have a broadcastable event_shape with
  ``num_steps = 1``, allowing :meth:`log_prob` to work with arbitrary length
  data::
  
      # homogeneous + homogeneous case:
      event_shape = (1,) + observation_dist.event_shape
  
  **References:**
  
  [1] Simo Sarkka, Angel F. Garcia-Fernandez (2019)
      "Temporal Parallelization of Bayesian Filters and Smoothers"
      https://arxiv.org/pdf/1905.13002.pdf
  
  :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_shape + (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_shape + (num_steps, state_dim, state_dim)``.
  :param ~torch.distributions.Distribution observation_dist: A conditional
      distribution of observed data conditioned on latent state. The
      ``.batch_shape`` should have rightmost size ``state_dim`` and be
      broadcastable to ``batch_shape + (num_steps, state_dim)``. The
      ``.event_shape`` may be arbitrary.
  :param int duration: Optional size of the time axis ``event_shape[0]``.
      This is required when sampling from homogeneous HMMs whose parameters
      are not expanded along the time axis.
class_docstr_tok: ["Hidden","Markov","Model","with","discrete","latent","state","and","arbitrary","observation","distribution",".","This","uses","[","1","]","to","parallelize","over","time",",","achieving","O","(","log","(","time",")",")","parallel","complexity","for","computing",":","meth",":","`","log_prob","`",",",":","meth",":","`","filter","`",",","and",":","meth",":","`","sample","`",".","The","event_shape","of","this","distribution","includes","time","on","the","left",":",":","event_shape","=","(","num_steps",",",")","+","observation_dist.event_shape","This","distribution","supports","any","combination","of","homogeneous\/heterogeneous","time","dependency","of","``","transition_logits","``","and","``","observation_dist","``",".","However",",","because","time","is","included","in","this","distribution","'s","event_shape",",","the","homogeneous+homogeneous","case","will","have","a","broadcastable","event_shape","with","``","num_steps","=","1","``",",","allowing",":","meth",":","`","log_prob","`","to","work","with","arbitrary","length","data",":",":","#","homogeneous","+","homogeneous","case",":","event_shape","=","(","1",",",")","+","observation_dist.event_shape","*","*","References",":","*","*","[","1","]","Simo","Sarkka",",","Angel","F.","Garcia-Fernandez","(","2019",")","``","Temporal","Parallelization","of","Bayesian","Filters","and","Smoothers","''","https",":","\/\/arxiv.org\/pdf\/1905.13002.pdf",":","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_shape","+","(","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_shape","+","(","num_steps",",","state_dim",",","state_dim",")","``",".",":","param","~torch.distributions.Distribution","observation_dist",":","A","conditional","distribution","of","observed","data","conditioned","on","latent","state",".","The","``",".batch_shape","``","should","have","rightmost","size","``","state_dim","``","and","be","broadcastable","to","``","batch_shape","+","(","num_steps",",","state_dim",")","``",".","The","``",".event_shape","``","may","be","arbitrary",".",":","param","int","duration",":","Optional","size","of","the","time","axis","``","event_shape","[","0","]","``",".","This","is","required","when","sampling","from","homogeneous","HMMs","whose","parameters","are","not","expanded","along","the","time","axis","."]
code: def __init__(
      self,
      initial_logits,
      transition_logits,
      observation_dist,
      validate_args=None,
      duration=None,
  ):
      if initial_logits.dim() < 1:
          raise ValueError(
              "expected initial_logits to have at least one dim, "
              "actual shape = {}".format(initial_logits.shape)
          )
      if transition_logits.dim() < 2:
          raise ValueError(
              "expected transition_logits to have at least two dims, "
              "actual shape = {}".format(transition_logits.shape)
          )
      if len(observation_dist.batch_shape) < 1:
          raise ValueError(
              "expected observation_dist to have at least one batch dim, "
              "actual .batch_shape = {}".format(
                  observation_dist.batch_shape
              )
          )
      shape = broadcast_shape(
          initial_logits.shape[:-1] + (1,),
          transition_logits.shape[:-2],
          observation_dist.batch_shape[:-1],
      )
      batch_shape, time_shape = shape[:-1], shape[-1:]
      event_shape = time_shape + observation_dist.event_shape
      self.initial_logits = initial_logits - initial_logits.logsumexp(
          -1, True
      )
      self.transition_logits = (
          transition_logits - transition_logits.logsumexp(-1, True)
      )
      self.observation_dist = observation_dist
      super().__init__(
          duration,
          batch_shape,
          event_shape,
          validate_args=validate_args,
      )
  
imports: import torch
  import torch.nn.functional
  from pyro.ops.gamma_gaussian import GammaGaussian, gamma_and_mvn_to_gamma_gaussian, gamma_gaussian_tensordot, matrix_and_mvn_to_gamma_gaussian
  from pyro.ops.gaussian import Gaussian, gaussian_tensordot, matrix_and_mvn_to_gaussian, mvn_to_gaussian, sequential_gaussian_filter_sample, sequential_gaussian_tensordot
  from pyro.ops.indexing import Vindex
  from pyro.ops.special import safe_log
  from pyro.ops.tensor_utils import cholesky_solve, safe_cholesky
  from .None import constraints
  from .torch import Categorical, Gamma, Independent, MultivariateNormal
  from .torch_distribution import TorchDistribution
  from .util import broadcast_shape, torch_jit_script_if_tracing
label_desc: Use image node_id 1 to create a new DiscreteHMM object from inherited base classes: HiddenMarkovModel with example: obj = DiscreteHMM(initial_logits, transition_logits, observation_dist, validate_args, duration)