code: def test_studentt_hmm_shape(
      batch_shape, duration, hidden_dim, obs_dim
  ):
      init_dist = random_studentt(batch_shape + (hidden_dim,)).to_event(
          1
      )
      trans_mat = torch.randn(
          batch_shape + (duration, hidden_dim, hidden_dim)
      )
      trans_dist = random_studentt(
          batch_shape + (duration, hidden_dim)
      ).to_event(1)
      obs_mat = torch.randn(
          batch_shape + (duration, hidden_dim, obs_dim)
      )
      obs_dist = random_studentt(
          batch_shape + (duration, obs_dim)
      ).to_event(1)
      hmm = dist.LinearHMM(
          init_dist,
          trans_mat,
          trans_dist,
          obs_mat,
          obs_dist,
          duration=duration,
      )
  
      def model(data=None):
          with pyro.plate_stack("plates", batch_shape):
              return pyro.sample("x", hmm, obs=data)
  
      data = model()
      rep = StudentTReparam()
      with poutine.trace() as tr:
          with poutine.reparam(
              config={"x": LinearHMMReparam(rep, rep, rep)}
          ):
              model(data)
  
      assert isinstance(tr.trace.nodes["x"]["fn"], dist.GaussianHMM)
      assert tr.trace.nodes["x_init_gamma"]["fn"].event_shape == (
          hidden_dim,
      )
      assert tr.trace.nodes["x_trans_gamma"]["fn"].event_shape == (
          duration,
          hidden_dim,
      )
      assert tr.trace.nodes["x_obs_gamma"]["fn"].event_shape == (
          duration,
          obs_dim,
      )
      tr.trace.compute_log_prob()
  
file_path: pyro/tests/infer/reparam/test_hmm.py
imports: import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.infer.reparam import LinearHMMReparam, StableReparam, StudentTReparam, SymmetricStableReparam
  from tests.common import assert_close
  from tests.ops.gaussian import random_mvn
  from .util import check_init_reparam
label_desc: Use image node_id 4 for calling a global function with example usage: test_studentt_hmm_shape(batch_shape, duration, hidden_dim, obs_dim) and returns: pyro