code: def test_stable_hmm_smoke(
      batch_shape, num_steps, hidden_dim, obs_dim
  ):
      init_dist = random_stable(batch_shape + (hidden_dim,)).to_event(1)
      trans_mat = torch.randn(
          batch_shape + (num_steps, hidden_dim, hidden_dim),
          requires_grad=True,
      )
      trans_dist = random_stable(
          batch_shape + (num_steps, hidden_dim)
      ).to_event(1)
      obs_mat = torch.randn(
          batch_shape + (num_steps, hidden_dim, obs_dim),
          requires_grad=True,
      )
      obs_dist = random_stable(
          batch_shape + (num_steps, obs_dim)
      ).to_event(1)
      data = obs_dist.sample()
      assert data.shape == batch_shape + (num_steps, obs_dim)
  
      def model(data):
          hmm = dist.LinearHMM(
              init_dist,
              trans_mat,
              trans_dist,
              obs_mat,
              obs_dist,
              duration=num_steps,
          )
          with pyro.plate_stack("plates", batch_shape):
              z = pyro.sample("z", hmm)
              pyro.sample("x", dist.Normal(z, 1).to_event(2), obs=data)
  
      # Test that we can combine these two reparameterizers.
      reparam_model = poutine.reparam(
          model,
          {
              "z": LinearHMMReparam(
                  StableReparam(), StableReparam(), StableReparam()
              ),
          },
      )
      reparam_model = poutine.reparam(
          reparam_model,
          {
              "z": ConjugateReparam(dist.Normal(data, 1).to_event(2)),
          },
      )
      reparam_guide = AutoDiagonalNormal(
          reparam_model
      )  # Models auxiliary variables.
  
      # Smoke test only.
      elbo = Trace_ELBO(num_particles=5, vectorize_particles=True)
      loss = elbo.differentiable_loss(
          reparam_model, reparam_guide, data
      )
      params = [trans_mat, obs_mat]
      torch.autograd.grad(loss, params, retain_graph=True)
  
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 6 for calling a global function with example usage: test_stable_hmm_smoke(batch_shape, num_steps, hidden_dim, obs_dim) without return types