code: def test_gaussian_hmm_elbo(
      batch_shape, num_steps, hidden_dim, obs_dim
  ):
      init_dist = random_mvn(batch_shape, hidden_dim)
      trans_mat = torch.randn(
          batch_shape + (num_steps, hidden_dim, hidden_dim),
          requires_grad=True,
      )
      trans_dist = random_mvn(batch_shape + (num_steps,), hidden_dim)
      obs_mat = torch.randn(
          batch_shape + (num_steps, hidden_dim, obs_dim),
          requires_grad=True,
      )
      obs_dist = random_mvn(batch_shape + (num_steps,), obs_dim)
  
      data = obs_dist.sample()
      assert data.shape == batch_shape + (num_steps, obs_dim)
      prior = dist.GaussianHMM(
          init_dist, trans_mat, trans_dist, obs_mat, obs_dist
      )
      likelihood = dist.Normal(data, 1).to_event(2)
      posterior, log_normalizer = prior.conjugate_update(likelihood)
  
      def model(data):
          with pyro.plate_stack("plates", batch_shape):
              z = pyro.sample("z", prior)
              pyro.sample("x", dist.Normal(z, 1).to_event(2), obs=data)
  
      def guide(data):
          with pyro.plate_stack("plates", batch_shape):
              pyro.sample("z", posterior)
  
      reparam_model = poutine.reparam(
          model, {"z": ConjugateReparam(likelihood)}
      )
  
      def reparam_guide(data):
          pass
  
      elbo = Trace_ELBO(num_particles=1000, vectorize_particles=True)
      expected_loss = elbo.differentiable_loss(model, guide, data)
      actual_loss = elbo.differentiable_loss(
          reparam_model, reparam_guide, data
      )
      assert_close(actual_loss, expected_loss, atol=0.01)
  
      params = [trans_mat, obs_mat]
      expected_grads = torch.autograd.grad(
          expected_loss, params, retain_graph=True
      )
      actual_grads = torch.autograd.grad(
          actual_loss, params, retain_graph=True
      )
      for a, e in zip(actual_grads, expected_grads):
          assert_close(a, e, rtol=0.01)
  
imports: import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer import Predictive, Trace_ELBO
  from pyro.infer.autoguide import AutoDiagonalNormal
  from pyro.infer.mcmc.api import MCMC
  from pyro.infer.mcmc.hmc import HMC
  from pyro.infer.reparam import ConjugateReparam, LinearHMMReparam, StableReparam
  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_gaussian_hmm_elbo(batch_shape, num_steps, hidden_dim, obs_dim) without return types