code: def test_profile_alternate_imp(M, batch_size, substitute):
      # --- Setup random model. ---
      pf_arranger = Profile(M)
  
      u1 = torch.rand((M + 1, 3))
      u1[
          M, :
      ] = 0  # Assume u_{M+1, j} = 0 for j in {0, 1, 2} in Eqn. S40.
      u = torch.cat([(1 - u1)[:, :, None], u1[:, :, None]], dim=2)
      r1 = torch.rand((M + 1, 3))
      r1[
          M, :
      ] = 1  # Assume r_{M+1, j} = 1 for j in {0, 1, 2} in Eqn. S40.
      r = torch.cat([(1 - r1)[:, :, None], r1[:, :, None]], dim=2)
      s = torch.rand((M, 4))
      s = s / torch.sum(s, dim=1, keepdim=True)
      c = torch.rand((M + 1, 4))
      c = c / torch.sum(c, dim=1, keepdim=True)
  
      if batch_size is not None:
          s = torch.rand((batch_size, M, 4))
          s = s / torch.sum(s, dim=2, keepdim=True)
          u1 = torch.rand((batch_size, M + 1, 3))
          u1[:, M, :] = 0
          u = torch.cat(
              [(1 - u1)[:, :, :, None], u1[:, :, :, None]], dim=3
          )
  
      # Compute forward pass of state arranger to get HMM parameters.
      # Don't use dimension M, assumed fixed by statearranger.
      if substitute:
          ll = torch.rand((4, 5))
          ll = ll / torch.sum(ll, dim=1, keepdim=True)
          a0ln, aln, eln = pf_arranger.forward(
              torch.log(s),
              torch.log(c),
              torch.log(r[:-1, :]),
              torch.log(u[..., :-1, :, :]),
              torch.log(ll),
          )
      else:
          a0ln, aln, eln = pf_arranger.forward(
              torch.log(s),
              torch.log(c),
              torch.log(r[:-1, :]),
              torch.log(u[..., :-1, :, :]),
          )
  
      # - Remake HMM parameters to check. -
      # Here we implement Equation S40 from the MuE paper
      # (https://www.biorxiv.org/content/10.1101/2020.07.31.231381v1.full.pdf)
      # more directly, iterating over all the indices of the transition matrix
      # and initial transition vector.
      K = 2 * M + 1
      if batch_size is None:
          batch_dim_size = 1
          r1 = r1.unsqueeze(0)
          u1 = u1.unsqueeze(0)
          s = s.unsqueeze(0)
          c = c.unsqueeze(0)
          if substitute:
              ll = ll.unsqueeze(0)
      else:
          batch_dim_size = batch_size
          r1 = r1[None, :, :] * torch.ones([batch_size, 1, 1])
          c = c[None, :, :] * torch.ones([batch_size, 1, 1])
          if substitute:
              ll = ll.unsqueeze(0)
      expected_a = torch.zeros((batch_dim_size, K, K))
      expected_a0 = torch.zeros((batch_dim_size, K))
      expected_e = torch.zeros((batch_dim_size, K, 4))
      for b in range(batch_dim_size):
          m, g = -1, 0
          u1[b][-1] = 1e-32
          for gp in range(2):
              for mp in range(M + gp):
                  kp = mg2k(mp, gp, M)
                  if m + 1 - g == mp and gp == 0:
                      expected_a0[b, kp] = (1 - r1[b, m + 1 - g, g]) * (
                          1 - u1[b, m + 1 - g, g]
                      )
                  elif m + 1 - g < mp and gp == 0:
                      expected_a0[b, kp] = (
                          (1 - r1[b, m + 1 - g, g])
                          * u1[b, m + 1 - g, g]
                          * simpleprod(
                              [
                                  (1 - r1[b, mpp, 2]) * u1[b, mpp, 2]
                                  for mpp in range(m + 2 - g, mp)
                              ]
                          )
                          * (1 - r1[b, mp, 2])
                          * (1 - u1[b, mp, 2])
                      )
                  elif m + 1 - g == mp and gp == 1:
                      expected_a0[b, kp] = r1[b, m + 1 - g, g]
                  elif m + 1 - g < mp and gp == 1:
                      expected_a0[b, kp] = (
                          (1 - r1[b, m + 1 - g, g])
                          * u1[b, m + 1 - g, g]
                          * simpleprod(
                              [
                                  (1 - r1[b, mpp, 2]) * u1[b, mpp, 2]
                                  for mpp in range(m + 2 - g, mp)
                              ]
                          )
                          * r1[b, mp, 2]
                      )
          for g in range(2):
              for m in range(M + g):
                  k = mg2k(m, g, M)
                  for gp in range(2):
                      for mp in range(M + gp):
                          kp = mg2k(mp, gp, M)
                          if m + 1 - g == mp and gp == 0:
                              expected_a[b, k, kp] = (
                                  1 - r1[b, m + 1 - g, g]
                              ) * (1 - u1[b, m + 1 - g, g])
                          elif m + 1 - g < mp and gp == 0:
                              expected_a[b, k, kp] = (
                                  (1 - r1[b, m + 1 - g, g])
                                  * u1[b, m + 1 - g, g]
                                  * simpleprod(
                                      [
                                          (1 - r1[b, mpp, 2])
                                          * u1[b, mpp, 2]
                                          for mpp in range(
                                              m + 2 - g, mp
                                          )
                                      ]
                                  )
                                  * (1 - r1[b, mp, 2])
                                  * (1 - u1[b, mp, 2])
                              )
                          elif m + 1 - g == mp and gp == 1:
                              expected_a[b, k, kp] = r1[b, m + 1 - g, g]
                          elif m + 1 - g < mp and gp == 1:
                              expected_a[b, k, kp] = (
                                  (1 - r1[b, m + 1 - g, g])
                                  * u1[b, m + 1 - g, g]
                                  * simpleprod(
                                      [
                                          (1 - r1[b, mpp, 2])
                                          * u1[b, mpp, 2]
                                          for mpp in range(
                                              m + 2 - g, mp
                                          )
                                      ]
                                  )
                                  * r1[b, mp, 2]
                              )
                          elif (
                              m == M and mp == M and g == 0 and gp == 0
                          ):
                              expected_a[b, k, kp] = 1.0
  
          for g in range(2):
              for m in range(M + g):
                  k = mg2k(m, g, M)
                  if g == 0:
                      expected_e[b, k, :] = s[b, m, :]
                  else:
                      expected_e[b, k, :] = c[b, m, :]
      if substitute:
          expected_e = torch.matmul(expected_e, ll)
  
      # --- Check ---
      if batch_size is None:
          expected_a = expected_a.squeeze()
          expected_a0 = expected_a0.squeeze()
          expected_e = expected_e.squeeze()
  
          assert torch.allclose(
              torch.sum(torch.exp(a0ln)),
              torch.tensor(1.0),
              atol=1e-3,
              rtol=1e-3,
          )
          assert torch.allclose(
              torch.sum(torch.exp(aln), axis=1),
              torch.ones(2 * M + 1),
              atol=1e-3,
              rtol=1e-3,
          )
      assert torch.allclose(expected_a0, torch.exp(a0ln))
      assert torch.allclose(expected_a, torch.exp(aln))
      assert torch.allclose(expected_e, torch.exp(eln))
  
imports: import pytest
  import torch
  from pyro.contrib.mue.statearrangers import Profile, mg2k
label_desc: Use image node_id 2 for calling a global function with example usage: test_profile_alternate_imp(M, batch_size, substitute) without return types