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 expand(self, batch_shape, _instance=None):
      new = self._get_checked_instance(DiscreteHMM, _instance)
      batch_shape = torch.Size(
          broadcast_shape(self.batch_shape, batch_shape)
      )
      # We only need to expand one of the inputs, since batch_shape is determined
      # by broadcasting all three. To save computation in _sequential_logmatmulexp(),
      # we expand only initial_logits, which is applied only after the logmatmulexp.
      # This is similar to the ._unbroadcasted_* pattern used elsewhere in distributions.
      new.initial_logits = self.initial_logits.expand(
          batch_shape + (-1,)
      )
      new.transition_logits = self.transition_logits
      new.observation_dist = self.observation_dist
      super(DiscreteHMM, new).__init__(
          self.duration,
          batch_shape,
          self.event_shape,
          validate_args=False,
      )
      new._validate_args = self.__dict__.get("_validate_args")
      return new
  
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 3 for calling the DiscreteHMM obj's underlying member method code with example usage: obj.expand(batch_shape, _instance) and returns: new