code: def test_sequential_logmatmulexp(batch_shape, state_dim, num_steps):
      logits = torch.randn(
          batch_shape + (num_steps, state_dim, state_dim)
      )
      actual = _sequential_logmatmulexp(logits)
      assert actual.shape == batch_shape + (state_dim, state_dim)
  
      # Check against einsum.
      operands = list(logits.unbind(-3))
      symbol = (opt_einsum.get_symbol(i) for i in range(1000))
      batch_symbols = "".join(next(symbol) for _ in batch_shape)
      state_symbols = [next(symbol) for _ in range(num_steps + 1)]
      equation = (
          ",".join(
              batch_symbols + state_symbols[t] + state_symbols[t + 1]
              for t in range(num_steps)
          )
          + "->"
          + batch_symbols
          + state_symbols[0]
          + state_symbols[-1]
      )
      expected = opt_einsum.contract(
          equation, *operands, backend="pyro.ops.einsum.torch_log"
      )
      assert_close(actual, expected)
  
imports: import logging
  import operator
  from functools import reduce
  import opt_einsum
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.distributions.hmm import _sequential_gamma_gaussian_tensordot, _sequential_logmatmulexp
  from pyro.distributions.util import broadcast_shape
  from pyro.infer import TraceEnum_ELBO, config_enumerate
  from pyro.ops.gamma_gaussian import gamma_and_mvn_to_gamma_gaussian, gamma_gaussian_tensordot, matrix_and_mvn_to_gamma_gaussian
  from pyro.ops.gaussian import gaussian_tensordot, matrix_and_mvn_to_gaussian, mvn_to_gaussian
  from pyro.ops.indexing import Vindex
  from tests.common import assert_close
  from tests.ops.gamma_gaussian import assert_close_gamma_gaussian, random_gamma, random_gamma_gaussian
  from tests.ops.gaussian import random_mvn
label_desc: Use image node_id 3 for calling a global function with example usage: test_sequential_logmatmulexp(batch_shape, state_dim, num_steps) without return types