class_docstr: FactorMuE
  
  This model consists of probabilistic PCA plus a MuE output distribution.
  
  The priors are all Normal distributions, and where relevant pushed through
  a softmax onto the simplex.
  
  :param int data_length: Length of the input sequence matrix, including
      zero padding at the end.
  :param int alphabet_length: Length of the sequence alphabet (e.g. 20 for
      amino acids).
  :param int z_dim: Number of dimensions of the z space.
  :param int batch_size: Minibatch size.
  :param int latent_seq_length: Length of the latent regressor sequence (M).
      Must be greater than or equal to 1. (Default: 1.1 x data_length.)
  :param bool indel_factor_dependence: Indel probabilities depend on the
      latent variable z.
  :param float indel_prior_scale: Standard deviation of the prior
      distribution on indel parameters.
  :param float indel_prior_bias: Mean of the prior distribution over the
      log probability of an indel not occurring. Higher values lead to lower
      probability of indels.
  :param float inverse_temp_prior: Mean of the prior distribution over the
      inverse temperature parameter.
  :param float weights_prior_scale: Standard deviation of the prior
      distribution over the factors.
  :param float offset_prior_scale: Standard deviation of the prior
      distribution over the offset (constant) in the pPCA model.
  :param str z_prior_distribution: Prior distribution over the latent
      variable z. Either 'Normal' (pPCA model) or 'Laplace' (an ICA model).
  :param bool ARD_prior: Use automatic relevance determination prior on
      factors.
  :param bool substitution_matrix: Use a learnable substitution matrix
      rather than the identity matrix.
  :param float substitution_prior_scale: Standard deviation of the prior
      distribution over substitution matrix parameters (when
      substitution_matrix is True).
  :param int latent_alphabet_length: Length of the alphabet in the latent
      regressor sequence.
  :param bool cuda: Transfer data onto the GPU during training.
  :param bool pin_memory: Pin memory for faster GPU transfer.
  :param float epsilon: A small value for numerical stability.
class_docstr_tok: ["FactorMuE","This","model","consists","of","probabilistic","PCA","plus","a","MuE","output","distribution",".","The","priors","are","all","Normal","distributions",",","and","where","relevant","pushed","through","a","softmax","onto","the","simplex",".",":","param","int","data_length",":","Length","of","the","input","sequence","matrix",",","including","zero","padding","at","the","end",".",":","param","int","alphabet_length",":","Length","of","the","sequence","alphabet","(","e.g",".","20","for","amino","acids",")",".",":","param","int","z_dim",":","Number","of","dimensions","of","the","z","space",".",":","param","int","batch_size",":","Minibatch","size",".",":","param","int","latent_seq_length",":","Length","of","the","latent","regressor","sequence","(","M",")",".","Must","be","greater","than","or","equal","to","1",".","(","Default",":","1.1","x","data_length",".",")",":","param","bool","indel_factor_dependence",":","Indel","probabilities","depend","on","the","latent","variable","z.",":","param","float","indel_prior_scale",":","Standard","deviation","of","the","prior","distribution","on","indel","parameters",".",":","param","float","indel_prior_bias",":","Mean","of","the","prior","distribution","over","the","log","probability","of","an","indel","not","occurring",".","Higher","values","lead","to","lower","probability","of","indels",".",":","param","float","inverse_temp_prior",":","Mean","of","the","prior","distribution","over","the","inverse","temperature","parameter",".",":","param","float","weights_prior_scale",":","Standard","deviation","of","the","prior","distribution","over","the","factors",".",":","param","float","offset_prior_scale",":","Standard","deviation","of","the","prior","distribution","over","the","offset","(","constant",")","in","the","pPCA","model",".",":","param","str","z_prior_distribution",":","Prior","distribution","over","the","latent","variable","z",".","Either","'Normal","'","(","pPCA","model",")","or","'Laplace","'","(","an","ICA","model",")",".",":","param","bool","ARD_prior",":","Use","automatic","relevance","determination","prior","on","factors",".",":","param","bool","substitution_matrix",":","Use","a","learnable","substitution","matrix","rather","than","the","identity","matrix",".",":","param","float","substitution_prior_scale",":","Standard","deviation","of","the","prior","distribution","over","substitution","matrix","parameters","(","when","substitution_matrix","is","True",")",".",":","param","int","latent_alphabet_length",":","Length","of","the","alphabet","in","the","latent","regressor","sequence",".",":","param","bool","cuda",":","Transfer","data","onto","the","GPU","during","training",".",":","param","bool","pin_memory",":","Pin","memory","for","faster","GPU","transfer",".",":","param","float","epsilon",":","A","small","value","for","numerical","stability","."]
code: def model(self, seq_data, local_scale, local_prior_scale):
      # ARD prior.
      if self.ARD_prior:
          # Relevance factors
          alpha = pyro.sample(
              "alpha",
              dist.Gamma(
                  torch.ones(self.z_dim), torch.ones(self.z_dim)
              ).to_event(1),
          )
      else:
          alpha = torch.ones(self.z_dim)
  
      # Factor and offset.
      W = pyro.sample(
          "W",
          dist.Normal(
              torch.zeros([self.z_dim, self.total_factor_size]),
              torch.ones([self.z_dim, self.total_factor_size])
              * self.weights_prior_scale
              / (alpha[:, None] + self.epsilon),
          ).to_event(2),
      )
      B = pyro.sample(
          "B",
          dist.Normal(
              torch.zeros(self.total_factor_size),
              torch.ones(self.total_factor_size)
              * self.offset_prior_scale,
          ).to_event(1),
      )
  
      # Indel probabilities.
      if not self.indel_factor_dependence:
          insert = pyro.sample(
              "insert",
              dist.Normal(
                  self.indel_prior * torch.ones(self.indel_shape),
                  self.indel_prior_scale * torch.ones(self.indel_shape),
              ).to_event(3),
          )
          insert_logits = insert - insert.logsumexp(-1, True)
          delete = pyro.sample(
              "delete",
              dist.Normal(
                  self.indel_prior * torch.ones(self.indel_shape),
                  self.indel_prior_scale * torch.ones(self.indel_shape),
              ).to_event(3),
          )
          delete_logits = delete - delete.logsumexp(-1, True)
  
      # Inverse temperature.
      inverse_temp = pyro.sample(
          "inverse_temp",
          dist.Normal(self.inverse_temp_prior, torch.tensor(1.0)),
      )
  
      # Substitution matrix.
      if self.substitution_matrix:
          substitute = pyro.sample(
              "substitute",
              dist.Normal(
                  torch.zeros(
                      [
                          self.latent_alphabet_length,
                          self.alphabet_length,
                      ]
                  ),
                  self.substitution_prior_scale
                  * torch.ones(
                      [
                          self.latent_alphabet_length,
                          self.alphabet_length,
                      ]
                  ),
              ).to_event(2),
          )
  
      with pyro.plate("batch", seq_data.shape[0]):
          with poutine.scale(scale=local_scale):
              with poutine.scale(scale=local_prior_scale):
                  # Sample latent variable from prior.
                  if self.z_prior_distribution == "Normal":
                      z = pyro.sample(
                          "latent",
                          dist.Normal(
                              torch.zeros(self.z_dim),
                              torch.ones(self.z_dim),
                          ).to_event(1),
                      )
                  elif self.z_prior_distribution == "Laplace":
                      z = pyro.sample(
                          "latent",
                          dist.Laplace(
                              torch.zeros(self.z_dim),
                              torch.ones(self.z_dim),
                          ).to_event(1),
                      )
  
              # Decode latent sequence.
              decoded = self.decoder(z, W, B, inverse_temp)
              if self.indel_factor_dependence:
                  insert_logits = decoded["insert_logits"]
                  delete_logits = decoded["delete_logits"]
  
              # Construct HMM parameters.
              if self.substitution_matrix:
                  (
                      initial_logits,
                      transition_logits,
                      observation_logits,
                  ) = self.statearrange(
                      decoded["precursor_seq_logits"],
                      decoded["insert_seq_logits"],
                      insert_logits,
                      delete_logits,
                      substitute,
                  )
              else:
                  (
                      initial_logits,
                      transition_logits,
                      observation_logits,
                  ) = self.statearrange(
                      decoded["precursor_seq_logits"],
                      decoded["insert_seq_logits"],
                      insert_logits,
                      delete_logits,
                  )
              # Draw samples.
              pyro.sample(
                  "obs_seq",
                  MissingDataDiscreteHMM(
                      initial_logits,
                      transition_logits,
                      observation_logits,
                  ),
                  obs=seq_data,
              )
  
imports: import datetime
  import numpy
  import torch
  import torch.nn
  from torch.nn.functional import softplus
  from torch.optim import Adam
  from torch.utils.data import DataLoader
  import pyro
  import pyro.distributions
  from pyro import poutine
  from pyro.contrib.mue.missingdatahmm import MissingDataDiscreteHMM
  from pyro.contrib.mue.statearrangers import Profile
  from pyro.infer import SVI, JitTrace_ELBO, Trace_ELBO
  from pyro.optim import MultiStepLR
label_desc: Use image node_id 3 for calling the FactorMuE obj's underlying member method code with example usage: obj.model(seq_data, local_scale, local_prior_scale) without return types