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.
  
  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
  
  This class should be interchangeable with
  :class:`pyro.distributions.hmm.DiscreteHMM` .
  
  **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.
code: def __init__(
      self,
      initial_logits,
      transition_logits,
      observation_dist,
      validate_args=None,
  ):
      assert isinstance(initial_logits, torch.Tensor)
      assert isinstance(transition_logits, torch.Tensor)
      assert isinstance(
          observation_dist, torch.distributions.Distribution
      )
      assert initial_logits.dim() >= 1
      assert transition_logits.dim() >= 2
      assert len(observation_dist.batch_shape) >= 1
      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._has_rsample = observation_dist.has_rsample
  
      # Normalize.
      initial_logits = initial_logits - initial_logits.logsumexp(
          -1, True
      )
      transition_logits = (
          transition_logits - transition_logits.logsumexp(-1, True)
      )
  
      # Convert tensors and distributions to funsors.
      init = tensor_to_funsor(initial_logits, ("state",))
      trans = tensor_to_funsor(
          transition_logits, ("time", "state", "state(time=1)")
      )
      obs = dist_to_funsor(observation_dist, ("time", "state(time=1)"))
      dtype = obs.inputs["value"].dtype
  
      # Construct the joint funsor.
      with lazy:
          # TODO perform math here once sequential_sum_product has been
          #   implemented as a first-class funsor.
          funsor_dist = Variable(
              "value", obs.inputs["value"]
          )  # a bogus value
          # Until funsor_dist is defined, we save factors for hand-computation in .log_prob().
          self._init = init
          self._trans = trans
          self._obs = obs
  
      super(DiscreteHMM, self).__init__(
          funsor_dist, batch_shape, event_shape, dtype, validate_args
      )
  
imports: from collections import OrderedDict
  import torch
  import funsor.ops
  from funsor.domains import Bint, Reals
  from funsor.interpretations import eager, lazy, moment_matching
  from funsor.pyro.convert import dist_to_funsor, funsor_to_cat_and_mvn, funsor_to_tensor, matrix_and_mvn_to_funsor, mvn_to_funsor, tensor_to_funsor
  from funsor.pyro.distribution import FunsorDistribution
  from funsor.sum_product import MarkovProduct, naive_sequential_sum_product, sequential_sum_product
  from funsor.terms import Variable
  from funsor.util import broadcast_shape
label_desc: Use image node_id 1 to create a new DiscreteHMM object from inherited base classes: FunsorDistribution with example: obj = DiscreteHMM(initial_logits, transition_logits, observation_dist, validate_args)