code: def test_profile_shapes(
      batch_ancestor_seq,
      batch_insert_seq,
      batch_insert,
      batch_delete,
      batch_substitute,
  ):
      M, D, B = 5, 2, 3
      K = 2 * M + 1
      batch_size = 6
      pf_arranger = Profile(M)
      sln = torch.randn([batch_size] * batch_ancestor_seq + [M, D])
      sln = sln - sln.logsumexp(-1, True)
      cln = torch.randn([batch_size] * batch_insert_seq + [M + 1, D])
      cln = cln - cln.logsumexp(-1, True)
      rln = torch.randn([batch_size] * batch_insert + [M, 3, 2])
      rln = rln - rln.logsumexp(-1, True)
      uln = torch.randn([batch_size] * batch_delete + [M, 3, 2])
      uln = uln - uln.logsumexp(-1, True)
      lln = torch.randn([batch_size] * batch_substitute + [D, B])
      lln = lln - lln.logsumexp(-1, True)
      a0ln, aln, eln = pf_arranger.forward(sln, cln, rln, uln, lln)
  
      if all(
          [
              not batch_ancestor_seq,
              not batch_insert_seq,
              not batch_substitute,
          ]
      ):
          assert eln.shape == (K, B)
          assert torch.allclose(eln.logsumexp(-1), torch.zeros(K))
      else:
          assert eln.shape == (batch_size, K, B)
          assert torch.allclose(
              eln.logsumexp(-1), torch.zeros(batch_size, K)
          )
  
      if all([not batch_insert, not batch_delete]):
          assert a0ln.shape == (K,)
          assert torch.allclose(a0ln.logsumexp(-1), torch.zeros(1))
          assert aln.shape == (K, K)
          assert torch.allclose(aln.logsumexp(-1), torch.zeros(K))
      else:
          assert a0ln.shape == (batch_size, K)
          assert torch.allclose(
              a0ln.logsumexp(-1), torch.zeros(batch_size)
          )
          assert aln.shape == (batch_size, K, K)
          assert torch.allclose(
              aln.logsumexp(-1), torch.zeros((batch_size, K))
          )
  
imports: import pytest
  import torch
  from pyro.contrib.mue.statearrangers import Profile, mg2k
label_desc: Use image node_id 3 for calling a global function with example usage: test_profile_shapes(batch_ancestor_seq, batch_insert_seq, batch_insert, batch_delete, batch_substitute) without return types