class_docstr: Profile HMM state arrangement. Parameterizes an HMM according to
  Equation S40 in [1] (with r_{M+1,j} = 1 and u_{M+1,j} = 0
  for j in {0, 1, 2}). For further background on profile HMMs see [2].
  
  **References**
  
  [1] E. N. Weinstein, D. S. Marks (2021)
  "Generative probabilistic biological sequence models that account for
  mutational variability"
  https://www.biorxiv.org/content/10.1101/2020.07.31.231381v2.full.pdf
  
  [2] R. Durbin, S. R. Eddy, A. Krogh, and G. Mitchison (1998)
  "Biological sequence analysis: probabilistic models of proteins and nucleic
  acids"
  Cambridge university press
  
  :param M: Length of regressor sequence.
  :type M: int
  :param epsilon: A small value for numerical stability.
  :type epsilon: float
docstr: Assemble HMM parameters given profile parameters.
  
  :param ~torch.Tensor precursor_seq_logits: Regressor sequence
      *log(x)*. Should have rightmost dimension ``(M, D)`` and be
      broadcastable to ``(batch_size, M, D)``, where
      D is the latent alphabet size. Should be normalized to one along the
      final axis, i.e. ``precursor_seq_logits.logsumexp(-1) = zeros``.
  :param ~torch.Tensor insert_seq_logits: Insertion sequence *log(c)*.
      Should have rightmost dimension ``(M+1, D)`` and be broadcastable
      to ``(batch_size, M+1, D)``. Should be normalized
      along the final axis.
  :param ~torch.Tensor insert_logits: Insertion probabilities *log(r)*.
      Should have rightmost dimension ``(M, 3, 2)`` and be broadcastable
      to ``(batch_size, M, 3, 2)``. Should be normalized along the
      final axis.
  :param ~torch.Tensor delete_logits: Deletion probabilities *log(u)*.
      Should have rightmost dimension ``(M, 3, 2)`` and be broadcastable
      to ``(batch_size, M, 3, 2)``. Should be normalized along the
      final axis.
  :param ~torch.Tensor substitute_logits: Substitution probabilities
      *log(l)*. Should have rightmost dimension ``(D, B)``, where
      B is the alphabet size of the data, and broadcastable to
      ``(batch_size, D, B)``. Must be normalized along the
      final axis.
  :return: *initial_logits*, *transition_logits*, and
      *observation_logits*. These parameters can be used to directly
      initialize the MissingDataDiscreteHMM distribution.
  :rtype: ~torch.Tensor, ~torch.Tensor, ~torch.Tensor
docstr_tok: ["Assemble","HMM","parameters","given","profile","parameters",".",":","param","~torch.Tensor","precursor_seq_logits",":","Regressor","sequence","*","log","(","x",")","*",".","Should","have","rightmost","dimension","``","(","M",",","D",")","``","and","be","broadcastable","to","``","(","batch_size",",","M",",","D",")","``",",","where","D","is","the","latent","alphabet","size",".","Should","be","normalized","to","one","along","the","final","axis",",","i.e",".","``","precursor_seq_logits.logsumexp","(","-1",")","=","zeros","``",".",":","param","~torch.Tensor","insert_seq_logits",":","Insertion","sequence","*","log","(","c",")","*",".","Should","have","rightmost","dimension","``","(","M+1",",","D",")","``","and","be","broadcastable","to","``","(","batch_size",",","M+1",",","D",")","``",".","Should","be","normalized","along","the","final","axis",".",":","param","~torch.Tensor","insert_logits",":","Insertion","probabilities","*","log","(","r",")","*",".","Should","have","rightmost","dimension","``","(","M",",","3",",","2",")","``","and","be","broadcastable","to","``","(","batch_size",",","M",",","3",",","2",")","``",".","Should","be","normalized","along","the","final","axis",".",":","param","~torch.Tensor","delete_logits",":","Deletion","probabilities","*","log","(","u",")","*",".","Should","have","rightmost","dimension","``","(","M",",","3",",","2",")","``","and","be","broadcastable","to","``","(","batch_size",",","M",",","3",",","2",")","``",".","Should","be","normalized","along","the","final","axis",".",":","param","~torch.Tensor","substitute_logits",":","Substitution","probabilities","*","log","(","l",")","*",".","Should","have","rightmost","dimension","``","(","D",",","B",")","``",",","where","B","is","the","alphabet","size","of","the","data",",","and","broadcastable","to","``","(","batch_size",",","D",",","B",")","``",".","Must","be","normalized","along","the","final","axis",".",":","return",":","*","initial_logits","*",",","*","transition_logits","*",",","and","*","observation_logits","*",".","These","parameters","can","be","used","to","directly","initialize","the","MissingDataDiscreteHMM","distribution",".",":","rtype",":","~torch.Tensor",",","~torch.Tensor",",","~torch.Tensor"]
code: def forward(
      self,
      precursor_seq_logits,
      insert_seq_logits,
      insert_logits,
      delete_logits,
      substitute_logits=None,
  ):
      """
      Assemble HMM parameters given profile parameters.
  
      :param ~torch.Tensor precursor_seq_logits: Regressor sequence
          *log(x)*. Should have rightmost dimension ``(M, D)`` and be
          broadcastable to ``(batch_size, M, D)``, where
          D is the latent alphabet size. Should be normalized to one along the
          final axis, i.e. ``precursor_seq_logits.logsumexp(-1) = zeros``.
      :param ~torch.Tensor insert_seq_logits: Insertion sequence *log(c)*.
          Should have rightmost dimension ``(M+1, D)`` and be broadcastable
          to ``(batch_size, M+1, D)``. Should be normalized
          along the final axis.
      :param ~torch.Tensor insert_logits: Insertion probabilities *log(r)*.
          Should have rightmost dimension ``(M, 3, 2)`` and be broadcastable
          to ``(batch_size, M, 3, 2)``. Should be normalized along the
          final axis.
      :param ~torch.Tensor delete_logits: Deletion probabilities *log(u)*.
          Should have rightmost dimension ``(M, 3, 2)`` and be broadcastable
          to ``(batch_size, M, 3, 2)``. Should be normalized along the
          final axis.
      :param ~torch.Tensor substitute_logits: Substitution probabilities
          *log(l)*. Should have rightmost dimension ``(D, B)``, where
          B is the alphabet size of the data, and broadcastable to
          ``(batch_size, D, B)``. Must be normalized along the
          final axis.
      :return: *initial_logits*, *transition_logits*, and
          *observation_logits*. These parameters can be used to directly
          initialize the MissingDataDiscreteHMM distribution.
      :rtype: ~torch.Tensor, ~torch.Tensor, ~torch.Tensor
      """
      initial_logits = (
          torch.einsum(
              "...ijk,ijkl->...l", delete_logits, self.u_transf_0
          )
          + torch.einsum(
              "...ijk,ijkl->...l", insert_logits, self.r_transf_0
          )
          + (-1 / self.epsilon) * self.null_transf_0
      )
      transition_logits = (
          torch.einsum(
              "...ijk,ijklf->...lf", delete_logits, self.u_transf
          )
          + torch.einsum(
              "...ijk,ijklf->...lf", insert_logits, self.r_transf
          )
          + (-1 / self.epsilon) * self.null_transf
      )
      # Broadcasting for concatenation.
      if len(precursor_seq_logits.size()) > len(
          insert_seq_logits.size()
      ):
          insert_seq_logits = insert_seq_logits.unsqueeze(0).expand(
              [precursor_seq_logits.size()[0], -1, -1]
          )
      elif len(insert_seq_logits.size()) > len(
          precursor_seq_logits.size()
      ):
          precursor_seq_logits = precursor_seq_logits.unsqueeze(
              0
          ).expand([insert_seq_logits.size()[0], -1, -1])
      seq_logits = torch.cat(
          [precursor_seq_logits, insert_seq_logits], dim=-2
      )
  
      # Option to include the substitution matrix.
      if substitute_logits is not None:
          observation_logits = torch.logsumexp(
              seq_logits.unsqueeze(-1)
              + substitute_logits.unsqueeze(-3),
              dim=-2,
          )
      else:
          observation_logits = seq_logits
  
      return initial_logits, transition_logits, observation_logits
  
file_path: pyro/pyro/contrib/mue/statearrangers.py
imports: import torch
  import torch.nn
label_desc: Use image node_id 3 for calling the Profile obj's underlying member method code with example usage: obj.forward(precursor_seq_logits, insert_seq_logits, insert_logits, delete_logits, substitute_logits) and returns: initial_logits,  transition_logits,  observation_logits