code: def test_profile_trivial_cases(M):
      # Trivial case: indel probabability of zero. Expected value of
      # HMM should match ancestral sequence times substitution matrix.
  
      # --- Setup model. ---
      D, B = 2, 2
      batch_size = 5
      pf_arranger = Profile(M)
      sln = torch.randn([batch_size, M, D])
      sln = sln - sln.logsumexp(-1, True)
      cln = torch.randn([batch_size, M + 1, D])
      cln = cln - cln.logsumexp(-1, True)
      rln = torch.cat(
          [
              torch.zeros([M, 3, 1]),
              -1 / pf_arranger.epsilon * torch.ones([M, 3, 1]),
          ],
          axis=-1,
      )
      uln = torch.cat(
          [
              torch.zeros([M, 3, 1]),
              -1 / pf_arranger.epsilon * torch.ones([M, 3, 1]),
          ],
          axis=-1,
      )
      lln = torch.randn([D, B])
      lln = lln - lln.logsumexp(-1, True)
  
      a0ln, aln, eln = pf_arranger.forward(sln, cln, rln, uln, lln)
  
      # --- Compute expected value per step. ---
      Eyln = torch.zeros([batch_size, M, B])
      ai = a0ln
      for j in range(M):
          Eyln[:, j, :] = torch.logsumexp(
              ai.unsqueeze(-1) + eln, axis=-2
          )
          ai = torch.logsumexp(ai.unsqueeze(-1) + aln, axis=-2)
  
      print(aln.exp())
      no_indel = torch.logsumexp(
          sln.unsqueeze(-1) + lln.unsqueeze(-3), axis=-2
      )
      assert torch.allclose(Eyln, no_indel)
  
imports: import pytest
  import torch
  from pyro.contrib.mue.statearrangers import Profile, mg2k
label_desc: Use image node_id 4 for calling a global function with example usage: test_profile_trivial_cases(M) without return types