class_docstr: Switching Linear Dynamical System represented as a Hidden Markov Model.
  
  This corresponds to the generative model::
  
      z = Categorical(logits=initial_logits).sample()
      y = initial_mvn[z].sample()
      x = []
      for t in range(num_steps):
          z = Categorical(logits=transition_logits[t, z]).sample()
          y = y @ transition_matrix[t, z] + transition_mvn[t, z].sample()
          x.append(y @ observation_matrix[t, z] + observation_mvn[t, z].sample())
  
  Viewed as a dynamic Bayesian network::
  
      z[t-1] ----> z[t] ---> z[t+1]         Discrete latent class
         |  \       |  \       |   \
         | y[t-1] ----> y[t] ----> y[t+1]   Gaussian latent state
         |   /      |   /      |   /
         V  /       V  /       V  /
      x[t-1]       x[t]      x[t+1]         Gaussian observation
  
  Let ``class`` be the latent class, ``state`` be the latent multivariate
  normal state, and ``value`` be the observed multivariate normal value.
  
  :param ~torch.Tensor initial_logits: Represents ``p(class[0])``.
  :param ~torch.distributions.MultivariateNormal initial_mvn: Represents
      ``p(state[0] | class[0])``.
  :param ~torch.Tensor transition_logits: Represents
      ``p(class[t+1] | class[t])``.
  :param ~torch.Tensor transition_matrix:
  :param ~torch.distributions.MultivariateNormal transition_mvn: Together
      with ``transition_matrix``, this represents
      ``p(state[t], state[t+1] | class[t])``.
  :param ~torch.Tensor observation_matrix:
  :param ~torch.distributions.MultivariateNormal observation_mvn: Together
      with ``observation_matrix``, this represents
      ``p(value[t+1], state[t+1] | class[t+1])``.
  :param bool exact: If True, perform exact inference at cost exponential in
      ``num_steps``. If False, use a :func:`~funsor.terms.moment_matching`
      approximation and use parallel scan algorithm to reduce parallel
      complexity to logarithmic in ``num_steps``. Defaults to False.
code: def __init__(
      self,
      initial_logits,
      initial_mvn,
      transition_logits,
      transition_matrix,
      transition_mvn,
      observation_matrix,
      observation_mvn,
      exact=False,
      validate_args=None,
  ):
      assert isinstance(initial_logits, torch.Tensor)
      assert isinstance(
          initial_mvn, torch.distributions.MultivariateNormal
      )
      assert isinstance(transition_logits, torch.Tensor)
      assert isinstance(transition_matrix, torch.Tensor)
      assert isinstance(
          transition_mvn, torch.distributions.MultivariateNormal
      )
      assert isinstance(observation_matrix, torch.Tensor)
      assert isinstance(
          observation_mvn, torch.distributions.MultivariateNormal
      )
      hidden_cardinality = initial_logits.size(-1)
      hidden_dim, obs_dim = observation_matrix.shape[-2:]
      assert (
          obs_dim >= hidden_dim // 2
      ), "obs_dim must be at least half of hidden_dim"
      assert initial_mvn.event_shape[0] == hidden_dim
      assert transition_logits.size(-1) == hidden_cardinality
      assert transition_matrix.shape[-2:] == (hidden_dim, hidden_dim)
      assert transition_mvn.event_shape[0] == hidden_dim
      assert observation_mvn.event_shape[0] == obs_dim
      init_shape = broadcast_shape(
          initial_logits.shape, initial_mvn.batch_shape
      )
      shape = broadcast_shape(
          init_shape[:-1] + (1, init_shape[-1]),
          transition_logits.shape[:-1],
          transition_matrix.shape[:-2],
          transition_mvn.batch_shape,
          observation_matrix.shape[:-2],
          observation_mvn.batch_shape,
      )
      assert shape[-1] == hidden_cardinality
      batch_shape, time_shape = shape[:-2], shape[-2:-1]
      event_shape = time_shape + (obs_dim,)
  
      # 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, ("class",)
      ) + dist_to_funsor(initial_mvn, ("class",))(value="state")
      trans = tensor_to_funsor(
          transition_logits, ("time", "class", "class(time=1)")
      ) + matrix_and_mvn_to_funsor(
          transition_matrix,
          transition_mvn,
          ("time", "class(time=1)"),
          "state",
          "state(time=1)",
      )
      obs = matrix_and_mvn_to_funsor(
          observation_matrix,
          observation_mvn,
          ("time", "class(time=1)"),
          "state(time=1)",
          "value",
      )
      if "class(time=1)" not in set(trans.inputs).union(obs.inputs):
          raise ValueError(
              "neither transition nor observation depend on discrete state"
          )
      dtype = "real"
  
      # 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(SwitchingLinearHMM, self).__init__(
          funsor_dist, batch_shape, event_shape, dtype, validate_args
      )
      self.exact = exact
  
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 SwitchingLinearHMM object from inherited base classes: FunsorDistribution with example: obj = SwitchingLinearHMM(initial_logits, initial_mvn, transition_logits, transition_matrix, transition_mvn, observation_matrix, observation_mvn, exact, validate_args)