class_docstr: This PyTorch Module encapsulates the model as well as the
  variational distribution (the guide) for the Deep Markov Model
class_docstr_tok: ["This","PyTorch","Module","encapsulates","the","model","as","well","as","the","variational","distribution","(","the","guide",")","for","the","Deep","Markov","Model"]
code: def model(
      self,
      mini_batch,
      mini_batch_reversed,
      mini_batch_mask,
      mini_batch_seq_lengths,
      annealing_factor=1.0,
  ):
      # this is the number of time steps we need to process in the mini-batch
      T_max = mini_batch.size(1)
  
      # register all PyTorch (sub)modules with pyro
      # this needs to happen in both the model and guide
      pyro.module("dmm", self)
  
      # set z_prev = z_0 to setup the recursive conditioning in p(z_t | z_{t-1})
      z_prev = self.z_0.expand(mini_batch.size(0), self.z_0.size(0))
  
      # we enclose all the sample statements in the model in a plate.
      # this marks that each datapoint is conditionally independent of the others
      with pyro.plate("z_minibatch", len(mini_batch)):
          # sample the latents z and observed x's one time step at a time
          # we wrap this loop in pyro.markov so that TraceEnum_ELBO can use multiple samples from the guide at each z
          for t in pyro.markov(range(1, T_max + 1)):
              # the next chunk of code samples z_t ~ p(z_t | z_{t-1})
              # note that (both here and elsewhere) we use poutine.scale to take care
              # of KL annealing. we use the mask() method to deal with raggedness
              # in the observed data (i.e. different sequences in the mini-batch
              # have different lengths)
  
              # first compute the parameters of the diagonal gaussian distribution p(z_t | z_{t-1})
              z_loc, z_scale = self.trans(z_prev)
  
              # then sample z_t according to dist.Normal(z_loc, z_scale)
              # note that we use the reshape method so that the univariate Normal distribution
              # is treated as a multivariate Normal distribution with a diagonal covariance.
              with poutine.scale(scale=annealing_factor):
                  z_t = pyro.sample(
                      "z_%d" % t,
                      dist.Normal(z_loc, z_scale)
                      .mask(mini_batch_mask[:, t - 1 : t])
                      .to_event(1),
                  )
  
              # compute the probabilities that parameterize the bernoulli likelihood
              emission_probs_t = self.emitter(z_t)
              # the next statement instructs pyro to observe x_t according to the
              # bernoulli distribution p(x_t|z_t)
              pyro.sample(
                  "obs_x_%d" % t,
                  dist.Bernoulli(emission_probs_t)
                  .mask(mini_batch_mask[:, t - 1 : t])
                  .to_event(1),
                  obs=mini_batch[:, t - 1, :],
              )
              # the latent sampled at this time step will be conditioned upon
              # in the next time step so keep track of it
              z_prev = z_t
  
imports: import argparse
  import logging
  import time
  from os.path import exists
  import numpy
  import torch
  import torch.nn
  import pyro
  import pyro.contrib.examples.polyphonic_data_loader
  import pyro.distributions
  import pyro.poutine
  from pyro.distributions import TransformedDistribution
  from pyro.distributions.transforms import affine_autoregressive
  from pyro.infer import SVI, JitTrace_ELBO, Trace_ELBO, TraceEnum_ELBO, TraceTMC_ELBO, config_enumerate
  from pyro.optim import ClippedAdam
label_desc: Use image node_id 2 for calling the DMM obj's underlying member method code with example usage: obj.model(mini_batch, mini_batch_reversed, mini_batch_mask, mini_batch_seq_lengths, annealing_factor) without return types